Tidying clients

This commit is contained in:
Cyber MacGeddon 2025-04-19 13:41:33 +01:00
parent 93bb2adb9e
commit dc0bcba269
2 changed files with 10 additions and 9 deletions

View file

@ -4,12 +4,17 @@ from .. schema import EmbeddingsReqeust, EmbeddingsResponse
class EmbeddingsClient(RequestResponse): class EmbeddingsClient(RequestResponse):
async def embed(self, text, timeout=30): async def embed(self, text, timeout=30):
resp = await self.request( resp = await self.request(
EmbeddingsRequest( EmbeddingsRequest(
text = text text = text
), ),
timeout=timeout timeout=timeout
) )
if resp.error:
raise RuntimeError(resp.error.message)
return resp.vectors return resp.vectors
class EmbeddingsClientSpec(RequestResponseSpec): class EmbeddingsClientSpec(RequestResponseSpec):

View file

@ -13,7 +13,7 @@ from .... schema import PromptRequest, PromptResponse, Error
from .... schema import TextCompletionRequest, TextCompletionResponse from .... schema import TextCompletionRequest, TextCompletionResponse
from .... base import FlowProcessor from .... base import FlowProcessor
from .... base import ProducerSpec, ConsumerSpec, RequestResponseSpec from .... base import ProducerSpec, ConsumerSpec, TextCompletionSpec
from . prompt_manager import PromptConfiguration, Prompt, PromptManager from . prompt_manager import PromptConfiguration, Prompt, PromptManager
@ -43,11 +43,9 @@ class Processor(FlowProcessor):
) )
self.register_specification( self.register_specification(
RequestResponseSpec( TextCompletionSpec(
request_name = "text-completion-request", request_name = "text-completion-request",
request_schema = TextCompletionRequest,
response_name = "text-completion-response", response_name = "text-completion-response",
response_schema = TextCompletionResponse,
) )
) )
@ -139,14 +137,12 @@ class Processor(FlowProcessor):
print(system, flush=True) print(system, flush=True)
print(prompt, flush=True) print(prompt, flush=True)
resp = await flow("text-completion-request").request( resp = await flow("text-completion-request").text_completion(
TextCompletionRequest( system = system, prompt = prompt,
system=system, prompt=prompt
),
) )
try: try:
return resp.response return resp
except Exception as e: except Exception as e:
print("LLM Exception:", e, flush=True) print("LLM Exception:", e, flush=True)
return None return None