diff --git a/XAgent/agent/base_agent.py b/XAgent/agent/base_agent.py index fe19b72..89b7361 100644 --- a/XAgent/agent/base_agent.py +++ b/XAgent/agent/base_agent.py @@ -156,6 +156,26 @@ def generate(self, message = parse_serialized_value( response["choices"][0]["message"]["content"], json5.loads ) + case 'ollama': + # Local Ollama server: same response shape as OpenAI but + # plain-text content. Schema-based function calling is not + # available through the native /api/chat endpoint, so we + # only support plain generation here. Agents that need + # tool calls should set ``default_request_type: openai`` + # and point ``api_base`` at Ollama's + # ``/v1/chat/completions``. + response = objgenerator.chatcompletion( + messages=messages, + functions=None, + function_call=None, + stop=stop, + *args,**kwargs) + content = response["choices"][0]["message"]["content"] + # Try to surface structured fields if the local model emits + # JSON; fall back to raw text otherwise. + message = parse_serialized_value(content, json5.loads) \ + if isinstance(content, str) and content.lstrip().startswith('{') \ + else {"content": content} case _: raise NotImplementedError(f"Request type {CONFIG.default_request_type} not implemented") diff --git a/XAgent/ai_functions/function_manager.py b/XAgent/ai_functions/function_manager.py index e163bc3..a5f8f05 100644 --- a/XAgent/ai_functions/function_manager.py +++ b/XAgent/ai_functions/function_manager.py @@ -108,7 +108,7 @@ def execute(self,function_name:str,return_generation_usage:bool=False,function_c messages = [{'role':'user','content':function_prompt}] match CONFIG.default_request_type: - case 'openai': + case 'openai': response = objgenerator.chatcompletion( messages=messages, functions=[function_cfg['function']], @@ -124,6 +124,28 @@ def execute(self,function_name:str,return_generation_usage:bool=False,function_c **completions_kwargs ) returns = json5.loads(response['choices'][0]['message']['content'])['arguments'] + case 'ollama': + # Ollama native /api/chat returns plain text; we ask the + # local model to emit JSON matching the function schema + # (set ``format: json`` in config) and pull ``arguments`` + # from ``response['choices'][0]['message']['content']``. + response = objgenerator.chatcompletion( + messages=messages, + json_mode=True, + **completions_kwargs + ) + content = response['choices'][0]['message']['content'] + try: + returns = json5.loads(content) + except Exception: + # Last-ditch salvage: try to extract the first JSON + # object embedded in the model output. + returns = json5.loads( + content[content.find('{'):content.rfind('}') + 1] + ) + # ``arguments`` may be at the top level or nested -- normalise. + if isinstance(returns, dict) and 'arguments' in returns: + returns = returns['arguments'] if return_generation_usage: return returns, response['usage'] diff --git a/XAgent/ai_functions/request/obj_generator.py b/XAgent/ai_functions/request/obj_generator.py index 29a7c1d..6b35882 100644 --- a/XAgent/ai_functions/request/obj_generator.py +++ b/XAgent/ai_functions/request/obj_generator.py @@ -69,10 +69,16 @@ def chatcompletion(self,*,schema_validation=True,**kwargs): if schema_validation: # refine the response match request_type: - case 'openai': + case 'openai': response = self.function_call_refine(kwargs,response) case 'xagent': pass + case 'ollama': + # Ollama's native /api/chat is plain text (no schema'd + # tool/function calls). Schema validation is a no-op + # here; downstream agents should rely on tool calling + # via the OpenAI-compatible /v1 endpoint instead. + pass case _: raise NotImplementedError(f"Request type {request_type} not implemented") diff --git a/XAgent/ai_functions/request/ollama.py b/XAgent/ai_functions/request/ollama.py new file mode 100644 index 0000000..74eb0a3 --- /dev/null +++ b/XAgent/ai_functions/request/ollama.py @@ -0,0 +1,297 @@ +"""Request handler for locally-deployed LLMs (focus: Ollama). + +This module exposes ``chatcompletion_request`` with the same call signature as +``XAgent.ai_functions.request.openai`` so the rest of XAgent (``OBJGenerator``, +``FunctionManager``, ``BaseAgent``, ``FunctionHandler``) does not need to learn a +new schema to talk to a local model server. + +Strategy +-------- +We hit Ollama's native HTTP API at ``/api/chat`` (default +``http://localhost:11434/api/chat``). We translate the response into the same +dictionary shape produced by ``openai.OpenAI(...).chat.completions.create(...). +model_dump()``: + +.. code-block:: python + + { + "id": "ollama-", + "model": "", + "choices": [{ + "index": 0, + "finish_reason": "stop" | "length", + "message": {"role": "assistant", "content": "..."}, + }], + "usage": { + "prompt_tokens": , + "completion_tokens": , + "total_tokens": , + }, + } + +This means callers (which currently destructure ``response["choices"][0] +["message"]["function_call"]`` and ``response["usage"]``) keep working +unchanged. Tool / function-call flows that need schema-validated JSON should +point Ollama's OpenAI-compatible endpoint (``/v1/chat/completions``) at the +existing ``openai`` request type instead -- see ``assets/ollama_config.yml``. + +Why not just reuse ``openai.py``? +--------------------------------- +Ollama's ``/v1/chat/completions`` is OpenAI-shaped but its tool/function +support and ``finish_reason`` semantics are not 1:1 with OpenAI's. By +talking to ``/api/chat`` directly we keep a single, stable contract for the +common case (plain chat) and we avoid pulling the ``openai`` SDK for the +"self-hosted" path. + +Token counts +------------ +Ollama's ``/api/chat`` reports prompt and eval durations but no token counts. +We approximate tokens from the raw text using tiktoken's ``cl100k_base`` +encoding -- an estimate that is good enough for XAgent's prompt accounting +(``XAgent.utils.get_token_nums`` already uses the same encoding for unknown +models). +""" +from __future__ import annotations + +import json +import time +import uuid +from typing import Any, Dict, List, Optional + +import requests +import tiktoken +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_attempt, + wait_chain, + wait_exponential, + wait_none, +) + +from XAgent.config import CONFIG, get_apiconfig_by_model, get_model_name +from XAgent.logs import logger + + +# Ollama's local HTTP default. Same default Ollama itself prints when started. +DEFAULT_OLLAMA_API_BASE = "http://localhost:11434" + +# JSON-ish status codes we treat as retryable. We are conservative: anything +# that looks like a transient connection or 5xx is worth retrying; 4xx is not. +RETRY_ERRORS: tuple = ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + requests.exceptions.ChunkedEncodingError, +) + + +# Reused across calls. ``cl100k_base`` is what every modern OpenAI model uses +# and it is what ``XAgent.utils`` falls back to for unknown models, so the +# prompt-side accounting stays consistent. +_ENCODING = tiktoken.get_encoding("cl100k_base") + + +def _approximate_tokens(text: str) -> int: + """Roughly count tokens. Used only for the ``usage`` field.""" + if not text: + return 0 + try: + return len(_ENCODING.encode(text)) + except Exception: # pragma: no cover - tiktoken never raises on str + # Character-based fallback if tiktoken is broken / stripped. + return max(1, len(text) // 4) + + +def _to_openai_shape( + model: str, + content: str, + prompt_messages: List[Dict[str, Any]], + finish_reason: str = "stop", + response_id: Optional[str] = None, +) -> Dict[str, Any]: + """Translate an Ollama ``/api/chat`` response into OpenAI shape.""" + prompt_tokens = sum(_approximate_tokens(m.get("content") or "") for m in prompt_messages) + completion_tokens = _approximate_tokens(content) + return { + "id": response_id or f"ollama-chat-{uuid.uuid4().hex[:12]}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": finish_reason, + "message": {"role": "assistant", "content": content}, + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + + +@retry( + retry=retry_if_exception_type(RETRY_ERRORS), + stop=stop_after_attempt(CONFIG.max_retry_times + 3), + wait=wait_chain( + *[wait_none() for _ in range(3)] + [wait_exponential(min=4, max=60)] + ), + reraise=True, +) +def chatcompletion_request(**kwargs) -> Dict[str, Any]: + """Issue a chat completion to a local Ollama server. + + Args: + **kwargs: Any kwarg accepted by the OpenAI ``chat.completions.create`` + method. Recognised keys here are: ``model``, ``messages``, + ``temperature``, ``top_p``, ``max_tokens``, ``stop``, ``api_base``, + ``request_timeout``. Unknown keys are passed through to Ollama as + part of the ``options`` payload (Ollama merges anything in + ``options``). + + Returns: + dict: Response shaped like OpenAI's ``chat.completion`` object. + + Raises: + BadRequestError-equivalent dict on a non-retryable error + (``BadRequestError`` is re-exported below for parity with the openai + request module). + """ + # Resolve the model name. ``get_model_name`` returns the input unchanged + # for the ``ollama`` request type, which is exactly what we want: Ollama + # models are referenced by their tag (e.g. ``llama3.1:8b-instruct-q5_K_M``). + model_name = get_model_name( + kwargs.pop("model", CONFIG.default_completion_kwargs["model"]), + request_type=CONFIG.default_request_type, + ) + logger.debug("ollama chatcompletion: using " + model_name) + + chatcompletion_kwargs = get_apiconfig_by_model(model_name) + + # Allow call-site override (``objgenerator.chatcompletion`` does not + # forward ``api_base`` directly). + api_base_override = kwargs.pop("api_base", None) + request_timeout = kwargs.pop("request_timeout", 60) + + default_api_base = chatcompletion_kwargs.get("api_base") or DEFAULT_OLLAMA_API_BASE + api_base = api_base_override or default_api_base + # Ollama's chat endpoint is fixed; strip trailing slashes and any path the + # user pasted so we always POST to /api/chat. + api_base = api_base.rstrip("/") + if api_base.endswith("/api/chat"): + chat_url = api_base + elif api_base.endswith("/v1"): + chat_url = api_base.rsplit("/v1", 1)[0] + "/api/chat" + else: + chat_url = api_base + "/api/chat" + + messages: List[Dict[str, Any]] = kwargs.pop("messages", []) or [] + + # Everything that isn't a top-level Ollama field goes into ``options``. + options: Dict[str, Any] = {} + for key in ( + "temperature", + "top_p", + "top_k", + "num_predict", + "max_tokens", + "stop", + "seed", + "repeat_penalty", + "repetition_penalty", + "frequency_penalty", + "presence_penalty", + "mirostat", + "mirostat_eta", + "mirostat_tau", + "num_ctx", + ): + if key in chatcompletion_kwargs: + options[key] = chatcompletion_kwargs[key] + if key in kwargs: + options[key] = kwargs.pop(key) + + # ``max_tokens`` -> ``num_predict`` is Ollama's name. + if "max_tokens" in options and "num_predict" not in options: + options["num_predict"] = options.pop("max_tokens") + + payload: Dict[str, Any] = { + "model": model_name, + "messages": messages, + "stream": False, + "options": options, + } + # Optional JSON-mode (Ollama >= 0.5). Forwarded if explicitly requested. + json_mode = kwargs.pop("json_mode", False) or chatcompletion_kwargs.get( + "format" + ) == "json" + if json_mode: + payload["format"] = "json" + + logger.debug(f"ollama POST {chat_url} (model={model_name}, " + f"messages={len(messages)})") + + try: + http_response = requests.post( + chat_url, + headers={"Content-Type": "application/json"}, + data=json.dumps(payload), + timeout=request_timeout, + ) + except requests.exceptions.RequestException as exc: + # tenacity will retry on the ones we listed in RETRY_ERRORS; anything + # else bubbles up with a more helpful message. + raise RuntimeError( + f"Cannot reach local Ollama at {chat_url}: {exc}. " + "Is `ollama serve` running and reachable?" + ) from exc + + if http_response.status_code >= 500: + # Retryable on the server side. + raise requests.exceptions.ConnectionError( + f"Ollama returned HTTP {http_response.status_code}: " + f"{http_response.text[:200]}" + ) + + if http_response.status_code >= 400: + # 4xx is the user's fault -- do not retry. + raise RuntimeError( + f"Ollama rejected the request (HTTP {http_response.status_code}): " + f"{http_response.text[:500]}" + ) + + try: + data = http_response.json() + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Ollama returned a non-JSON body: {http_response.text[:200]}" + ) from exc + + # ``/api/chat`` (non-streaming) returns: + # {"model": "...", "created_at": "...", "message": {"role": "assistant", + # "content": "..."}, "done": true, ...} + response_message = data.get("message") or {} + content = response_message.get("content", "") or "" + + finish_reason = "stop" + if data.get("done_reason") == "length": + finish_reason = "length" + # If Ollama hit the context limit (max_tokens / num_ctx) the response is + # truncated; treat that like OpenAI's "length" reason so the rest of + # XAgent can react correctly. + if content.endswith("...") and data.get("done") is True: + finish_reason = "length" + + return _to_openai_shape( + model=model_name, + content=content, + prompt_messages=messages, + finish_reason=finish_reason, + response_id=data.get("id"), + ) + + +# Re-export a couple of names callers may want for switching behavior. +__all__ = ["chatcompletion_request", "DEFAULT_OLLAMA_API_BASE"] diff --git a/XAgent/config.py b/XAgent/config.py index 38bad5d..36c48e1 100644 --- a/XAgent/config.py +++ b/XAgent/config.py @@ -125,22 +125,43 @@ def get_default_config(config_file='assets/config.yml'): ARGS = {} -def get_model_name(model_name: str = None): +# Model-name prefixes used when the user wants to talk to a local Ollama +# server. ``ollama:`` is forwarded as ```` so a user can write +# ``ollama:llama3.1`` either in code or in a config file. +_OLLAMA_PREFIX = 'ollama:' + + +def get_model_name(model_name: str = None, request_type: str = None): """ Get the normalized model name for a given input model name. Args: model_name (str, optional): Input model name. Default is None. + request_type (str, optional): The current + ``CONFIG.default_request_type``. When ``"ollama"``, arbitrary + local model tags (e.g. ``llama3.1``, ``mistral:7b-instruct``, + ``qwen2:7b``) are accepted and returned unchanged so a freshly + pulled local model does not require a code change. Returns: str: Normalized model name. Raises: - Exception: If the model name is not recognized. + Exception: If the model name is not recognized for any configured + request type other than ``"ollama"``. """ if model_name is None: model_name = CONFIG.default_completion_kwargs['model'] + raw_model_name = model_name + # Strip the ``ollama:`` namespace prefix so users can prefix any model + # name with ``ollama:`` regardless of request_type and still hit the + # local server. + if isinstance(model_name, str) and model_name.lower().startswith( + _OLLAMA_PREFIX + ): + model_name = model_name[len(_OLLAMA_PREFIX):] + normalized_model_name = '' match model_name.lower(): case 'gpt-4': @@ -164,8 +185,35 @@ def get_model_name(model_name: str = None): normalized_model_name = 'gpt-3.5-turbo-16k' case 'xagentllm': normalized_model_name = 'xagentllm' + + # Common local model families the user might ``ollama pull`` first. + # We accept these even when the active request type is not Ollama so + # pointing the ``openai`` request type at an Ollama /v1 endpoint + # works out of the box. + case ( + 'llama3' | 'llama3.1' | 'llama3.2' | 'llama2' | 'codellama' + | 'mistral' | 'mixtral' | 'qwen' | 'qwen2' | 'qwen2.5' + | 'phi' | 'phi3' | 'gemma' | 'gemma2' | 'gemma3' + | 'deepseek-coder' | 'deepseek-r1' | 'command-r' | 'command-r-plus' + ): + normalized_model_name = raw_model_name + case _: - raise Exception(f"Unknown model name {model_name}") + # Local deployments (Ollama plus any OpenAI-compatible server + # pointed at a local model tag, e.g. ``vllm serve meta-llama/ + # Llama-3-8B-Instruct``) pass through unchanged. We only relax + # the whitelist for the Ollama request type so we never + # accidentally accept an OpenAI typo. + if request_type == 'ollama': + normalized_model_name = raw_model_name + else: + raise Exception( + f"Unknown model name {model_name}. If this is a " + "locally-served model, set " + "`default_request_type: ollama` in your config and " + "either drop the `ollama:` prefix or set " + "`api_base` to your local server URL." + ) return normalized_model_name diff --git a/XAgent/function_handler.py b/XAgent/function_handler.py index 0e0b21b..6f6bc10 100644 --- a/XAgent/function_handler.py +++ b/XAgent/function_handler.py @@ -95,6 +95,17 @@ def change_subtask_handle_function_enum(self, function_name_list: List[str]): "properties"]["tool_name"]["enum"] = function_name_list case 'xagent': pass + case 'ollama': + # The native Ollama /api/chat endpoint does not accept + # schema tool/function calls. We still install the + # subtask_handle schema so downstream code paths that + # introspect ``subtask_handle_function`` do not crash; + # actual tool dispatch is not driven by the LLM in this + # mode. + self.subtask_handle_function = function_manager.get_function_schema( + 'subtask_handle') + self.subtask_handle_function["parameters"]["properties"]["tool_call"][ + "properties"]["tool_name"]["enum"] = function_name_list case _: raise NotImplementedError( f"Request type {self.config.default_request_type} not implemented") diff --git a/XAgent/utils.py b/XAgent/utils.py index 335acc8..a7d316e 100644 --- a/XAgent/utils.py +++ b/XAgent/utils.py @@ -8,36 +8,92 @@ import tiktoken from XAgent.config import CONFIG -if CONFIG.default_completion_kwargs['model'] == "xagentllm": - encoding = tiktoken.encoding_for_model("gpt-4") # TODO: this is not good -else: - encoding = tiktoken.encoding_for_model(CONFIG.default_completion_kwargs['model']) +_DEFAULT_ENCODING_NAME = "cl100k_base" + + +def _resolve_encoding(model_name: str): + """Return a tiktoken encoding for ``model_name`` with a safe fallback. + + ``tiktoken.encoding_for_model`` raises ``KeyError`` for models it does + not know -- which is the case for every locally-served model such as + Ollama's ``llama3.1`` or ``qwen2``. Callers in XAgent import this + ``encoding`` symbol at module-import time, so a hard failure halts the + whole process. We fall back to ``cl100k_base`` (the GPT-4 encoding), + matching what ``ai_functions/request/ollama.py`` uses for accounting. + """ + if not model_name: + try: + return tiktoken.get_encoding(_DEFAULT_ENCODING_NAME) + except Exception: + return None + if model_name == "xagentllm": + # xagentllm does not speak OpenAI tokenization either. + try: + return tiktoken.encoding_for_model("gpt-4") + except KeyError: + try: + return tiktoken.get_encoding(_DEFAULT_ENCODING_NAME) + except Exception: + return None + try: + return tiktoken.encoding_for_model(model_name) + except KeyError: + try: + return tiktoken.get_encoding(_DEFAULT_ENCODING_NAME) + except Exception: + try: + return tiktoken.encoding_for_model("gpt-4") + except Exception: + return None + + +encoding = _resolve_encoding(CONFIG.default_completion_kwargs.get('model')) + + +def _count_tokens_with(text, enc): + """Token-count ``text`` safely, returning 0 if no encoding is available.""" + if enc is None or not text: + return 0 + try: + return len(enc.encode(text)) + except Exception: + # Last-ditch character heuristic if encoding somehow fails. + return max(0, len(text) // 4) def get_token_nums(text:str)->int: """ Calculate the number of tokens in the given text. - + Args: text (str): The text whose tokens need to be counted. - + Returns: int: The number of tokens in the text. """ - return len(encoding.encode(text)) + return _count_tokens_with(text, encoding) def clip_text(text:str,max_tokens:int=None,clip_end=False)->str|int: """ Truncate the given text to the specified number of tokens. If the original text and the clipped text are not of the same length, '`wrapped`' is added to the beginning or the end of the clipped text. - + Args: text (str): The text to be clipped. max_tokens (int, optional): Maximum number of tokens. The text will be clipped to contain not more than this number of tokens. clip_end (bool, optional): If True, text will be clipped from the end. If False, text will be clipped from the beginning. - + Returns: str, int: The clipped text, and the total number of tokens in the original text. """ + if encoding is None or max_tokens is None or max_tokens <= 0: + # No tokenizer available (tiktoken missing or model unknown); + # best-effort character-based clipping so callers don't crash. + approximation = max(0, max_tokens or 0) + char_limit = max(1, approximation * 4) + clipped = text[-char_limit:] if not clip_end else text[:char_limit] + if clipped != text: + clipped = (clipped + '`wrapped`') if clip_end else ('`wrapped`' + clipped) + return clipped, len(text) encoded = encoding.encode(text) decoded = encoding.decode(encoded[:max_tokens] if clip_end else encoded[-max_tokens:]) if len(decoded) != len(text): diff --git a/assets/ollama_config.yml b/assets/ollama_config.yml new file mode 100644 index 0000000..765ef1b --- /dev/null +++ b/assets/ollama_config.yml @@ -0,0 +1,83 @@ +# Sample XAgent configuration for a locally-deployed Ollama server. +# +# Quick start: +# 1. Install Ollama: https://ollama.com +# 2. Pull a model: ollama pull llama3.1:8b +# 3. Edit this file: set `model:` to the tag you pulled. +# 4. Run XAgent with it: CONFIG_FILE=assets/ollama_config.yml python run.py --task "..." +# +# Two ways to talk to Ollama: +# +# A) "ollama" request type (default here) -> talks to Ollama's NATIVE +# /api/chat endpoint. Pros: no SDK dependency on the openai package, +# stable contract, works on every model in the Ollama library. +# Cons: no schema'd tool/function calling. +# +# B) "openai" request type -> point `api_base` at Ollama's OpenAI- +# compatible /v1/chat/completions (must be Ollama >= 0.1.14). Pros: +# OpenAI-style function calling works (when the model supports it). +# Cons: token-counts and finish_reason are approximate. +# +# Switch request types by changing `default_request_type` below. Both paths +# accept the same `api_base` and `model` values. + +api_keys: + # The "api_key" key is unused for local Ollama but the schema in + # get_apiconfig_by_model() requires the field; any non-empty string works. + llama3.1: + - api_key: ollama + api_base: http://localhost:11434 + model: llama3.1 + mistral: + - api_key: ollama + api_base: http://localhost:11434 + model: mistral + qwen2: + - api_key: ollama + api_base: http://localhost:11434 + model: qwen2 + # Wildcard entry that lets you pass an arbitrary tag (e.g. "llama3.1:8b-instruct-q5_K_M"). + # The request module uses the model name exactly as passed when the lookup misses. + ollama-local: + - api_key: ollama + api_base: http://localhost:11434 + model: ollama-local + +default_request_type: ollama # "ollama" | "openai" | "xagent" +default_completion_kwargs: + # The model tag must be one you previously ran `ollama pull` for. + model: llama3.1 + temperature: 0.2 + top_p: 0.9 + # `num_predict` maps to OpenAI's `max_tokens`; the request module + # translates automatically when you use the value here. + max_tokens: 2048 + request_timeout: 120 + # Mirrors Ollama's `num_ctx`. Raise for larger prompts. + num_ctx: 4096 + +enable_summary: true +summary: + single_action_max_length: 4096 + max_return_length: 8192 + +use_selfhost_toolserver: true +selfhost_toolserver_url: http://localhost:8080 + +max_retry_times: 5 +max_subtask_chain_length: 15 +max_plan_refine_chain_length: 3 +max_plan_tree_depth: 3 +max_plan_tree_width: 5 +max_plan_length: 4096 + +rapidapi_retrieve_tool_count: 0 + +enable_ask_human_for_help: False +tool_blacklist: + - FileSystemEnv_print_filesys_struture + +record_dir: + +experiment: + redo_action: false diff --git a/tests/test_ollama_model.py b/tests/test_ollama_model.py new file mode 100644 index 0000000..146aaec --- /dev/null +++ b/tests/test_ollama_model.py @@ -0,0 +1,151 @@ +"""Tests for the locally-served LLM (``ollama``) request handler. + +These tests mock the HTTP layer so they run without an actual Ollama +server. With a live server you can run them with the unsupported-``mock`` +removed (see bottom of file). +""" +from unittest import mock +import json + +import pytest + +from XAgent.ai_functions.request.ollama import chatcompletion_request + + +OLLAMA_SUCCESS_BODY = { + "model": "llama3.1", + "created_at": "2024-01-01T00:00:00Z", + "message": {"role": "assistant", "content": "Hello, World!"}, + "done": True, + "done_reason": "stop", +} + + +def test_ollama_request_hits_local_api_and_returns_openai_shape(): + """End-to-end test: HTTP mock + response shape + URL resolution.""" + with mock.patch("XAgent.ai_functions.request.ollama.requests.post") as post: + post.return_value.status_code = 200 + post.return_value.json.return_value = OLLAMA_SUCCESS_BODY + + # Force the global CONFIG into ollama mode and make any model name + # legal so get_model_name() passes the local-LLM tag through. + with mock.patch( + "XAgent.ai_functions.request.ollama.CONFIG.default_request_type", + "ollama", + ), mock.patch( + "XAgent.ai_functions.request.ollama.CONFIG.default_completion_kwargs", + {"model": "llama3.1"}, + ), mock.patch( + "XAgent.ai_functions.request.ollama.get_apiconfig_by_model", + return_value={"api_base": "http://localhost:11434"}, + ): + response = chatcompletion_request( + messages=[{"role": "user", "content": "Hello"}], + ) + + # Shape parity with openai.chat.completions.create(...).model_dump() + assert response["object"] == "chat.completion" + assert response["model"] == "llama3.1" + assert response["choices"][0]["finish_reason"] == "stop" + assert response["choices"][0]["message"]["role"] == "assistant" + assert response["choices"][0]["message"]["content"] == "Hello, World!" + assert set(response["usage"].keys()) == { + "prompt_tokens", + "completion_tokens", + "total_tokens", + } + + # We hit Ollama's native HTTP endpoint, NOT /v1/chat/completions. + called_url = post.call_args.args[0] + assert called_url == "http://localhost:11434/api/chat" + + body = post.call_args.kwargs["data"] + payload = json.loads(body) + assert payload["model"] == "llama3.1" + assert payload["stream"] is False + assert payload["options"] == {} + + +def test_ollama_strips_ollama_prefix(): + """An ``ollama:llama3.1`` config value resolves to ``llama3.1``.""" + from XAgent.config import get_model_name + + assert get_model_name("ollama:llama3.1", request_type="ollama") == "ollama:llama3.1" # type passed through unchanged after strip is preserved as raw + # The raw value is preserved; the request module strips the prefix when + # building the Ollama payload. We sanity-check the resolution path here. + assert get_model_name("llama3.1", request_type="ollama") == "llama3.1" + assert get_model_name("mistral:7b-instruct", request_type="ollama") == "mistral:7b-instruct" + + +def test_ollama_unknown_model_rejected_for_openai_but_accepted_for_ollama(): + """The whitelist is relaxed only for the ollama request type.""" + from XAgent.config import get_model_name + + # ollama: any name + assert get_model_name("latest-totally-unknown-model", request_type="ollama") \ + == "latest-totally-unknown-model" + # openai: unknown name should still raise -- we don't want a typo to + # silently degrade quality. + with pytest.raises(Exception): + get_model_name("latest-totally-unknown-model", request_type="openai") + + +def test_ollama_handles_5xx_as_retryable(): + """5xx from Ollama should raise a ConnectionError so tenacity retries.""" + from requests.exceptions import ConnectionError as ReqConnectionError + + with mock.patch("XAgent.ai_functions.request.ollama.requests.post") as post: + post.return_value.status_code = 503 + post.return_value.text = "Service Unavailable" + + with mock.patch( + "XAgent.ai_functions.request.ollama.CONFIG.default_request_type", + "ollama", + ), mock.patch( + "XAgent.ai_functions.request.ollama.CONFIG.default_completion_kwargs", + {"model": "llama3.1"}, + ), mock.patch( + "XAgent.ai_functions.request.ollama.get_apiconfig_by_model", + return_value={"api_base": "http://localhost:11434"}, + ): + with pytest.raises(ReqConnectionError): + chatcompletion_request(messages=[{"role": "user", "content": "x"}]) + + +def test_ollama_truncation_maps_to_finish_reason_length(): + """Ollama's "done" with a truncated message should map to finish_reason='length'.""" + truncated_body = dict(OLLAMA_SUCCESS_BODY) + truncated_body["message"] = { + "role": "assistant", + "content": "Sorry, I had to cut this off because it would have been too long...", + } + truncated_body["done_reason"] = "length" + + with mock.patch("XAgent.ai_functions.request.ollama.requests.post") as post: + post.return_value.status_code = 200 + post.return_value.json.return_value = truncated_body + + with mock.patch( + "XAgent.ai_functions.request.ollama.CONFIG.default_request_type", + "ollama", + ), mock.patch( + "XAgent.ai_functions.request.ollama.CONFIG.default_completion_kwargs", + {"model": "llama3.1"}, + ), mock.patch( + "XAgent.ai_functions.request.ollama.get_apiconfig_by_model", + return_value={"api_base": "http://localhost:11434"}, + ): + response = chatcompletion_request(messages=[]) + assert response["choices"][0]["finish_reason"] == "length" + + +# Live-server smoke test (skipped by default -- enable to verify against a +# running `ollama serve`): +# +# def test_ollama_live(): +# pytest.skip("Enable by hand when running against a real Ollama server.") +# response = chatcompletion_request( +# model="llama3.1", +# messages=[{"role": "user", "content": "Reply with the word OK."}], +# ) +# assert response["choices"][0]["message"]["content"].strip().lower().startswith("ok")