diff --git a/agent/core/llm.py b/agent/core/llm.py index a9512b3..06e1e65 100644 --- a/agent/core/llm.py +++ b/agent/core/llm.py @@ -21,7 +21,13 @@ def __init__(self, config: LLMConfig): self.config = config self.base_url = config.base_url.rstrip("/") - def _request(self, endpoint: str, data: dict, stream: bool = False) -> dict: + def _request( + self, + endpoint: str, + data: dict, + stream: bool = False, + timeout: Optional[float] = None, + ) -> dict: url = f"{self.base_url}{endpoint}" payload = json.dumps(data).encode("utf-8") req = urllib.request.Request( @@ -30,8 +36,9 @@ def _request(self, endpoint: str, data: dict, stream: bool = False) -> dict: headers={"Content-Type": "application/json"}, method="POST", ) + effective_timeout = timeout if timeout is not None else self.config.timeout try: - resp = urllib.request.urlopen(req, timeout=self.config.timeout) + resp = urllib.request.urlopen(req, timeout=effective_timeout) if stream: return resp # Return response object for streaming return json.loads(resp.read().decode("utf-8")) @@ -77,6 +84,7 @@ def chat( messages: list[dict], tools: Optional[list[dict]] = ..., stream: Literal[False] = ..., + timeout: Optional[float] = ..., ) -> dict: ... @overload @@ -86,6 +94,7 @@ def chat( tools: Optional[list[dict]] = ..., *, stream: Literal[True], + timeout: Optional[float] = ..., ) -> Generator[str, None, dict]: ... def chat( @@ -93,6 +102,7 @@ def chat( messages: list[dict], tools: Optional[list[dict]] = None, stream: bool = False, + timeout: Optional[float] = None, ) -> Union[dict, Generator[str, None, dict]]: """Send a chat completion request to Ollama. @@ -100,6 +110,7 @@ def chat( messages: List of message dicts with 'role' and 'content'. tools: Optional list of tool definitions for function calling. stream: Whether to stream the response. + timeout: Optional per-request timeout override in seconds. Returns: Response dict with 'message' containing 'role', 'content', @@ -119,14 +130,18 @@ def chat( data["tools"] = tools if stream: - return self._stream_chat(data) + return self._stream_chat(data, timeout=timeout) - result = self._request("/api/chat", data) + result = self._request("/api/chat", data, timeout=timeout) return result - def _stream_chat(self, data: dict) -> Generator[str, None, dict]: + def _stream_chat( + self, + data: dict, + timeout: Optional[float] = None, + ) -> Generator[str, None, dict]: """Stream chat response, yielding content chunks.""" - resp = self._request("/api/chat", data, stream=True) + resp = self._request("/api/chat", data, stream=True, timeout=timeout) full_response = {"message": {"role": "assistant", "content": ""}} for line in resp: @@ -149,8 +164,16 @@ def _stream_chat(self, data: dict) -> Generator[str, None, dict]: return full_response - def generate(self, prompt: str) -> str: - """Simple text generation without chat format.""" + def generate(self, prompt: str, timeout: Optional[float] = None) -> str: + """Simple text generation without chat format. + + Args: + prompt: The prompt to generate from. + timeout: Optional per-request timeout override in seconds. + + Returns: + Generated text response. + """ data = { "model": self.config.model, "prompt": prompt, @@ -160,16 +183,28 @@ def generate(self, prompt: str) -> str: "num_predict": self.config.max_tokens, }, } - result = self._request("/api/generate", data) + result = self._request("/api/generate", data, timeout=timeout) return result.get("response", "") - def pull_model(self, model_name: Optional[str] = None) -> bool: - """Pull/download a model from Ollama registry.""" + def pull_model( + self, + model_name: Optional[str] = None, + timeout: Optional[float] = None, + ) -> bool: + """Pull/download a model from Ollama registry. + + Args: + model_name: Model name to pull. Uses config model if not specified. + timeout: Optional per-request timeout override in seconds. + + Returns: + True if pull succeeded, False otherwise. + """ model = model_name or self.config.model print(f"Pulling model '{model}'... This may take a while.") try: data = {"name": model, "stream": False} - self._request("/api/pull", data) + self._request("/api/pull", data, timeout=timeout) print(f"Model '{model}' pulled successfully.") return True except OllamaError as e: