From cb907ea201226bea259fbe437a315617ca3c8535 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 30 Dec 2025 17:38:58 +0800 Subject: [PATCH 01/10] feat(utils): add universal timer decorator with loguru integration --- docs/deprecated.txt | 9 +++++ reme_ai/core/utils/__init__.py | 5 +++ reme_ai/core/utils/timer.py | 66 ++++++++++++++++++++++++++++++++++ tests/test_timer.py | 64 +++++++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+) create mode 100644 docs/deprecated.txt create mode 100644 reme_ai/core/utils/__init__.py create mode 100644 reme_ai/core/utils/timer.py create mode 100644 tests/test_timer.py diff --git a/docs/deprecated.txt b/docs/deprecated.txt new file mode 100644 index 00000000..41d9c205 --- /dev/null +++ b/docs/deprecated.txt @@ -0,0 +1,9 @@ +from loguru import logger + +用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码 +C0114: Missing module docstring (missing-module-docstring) +C0115: Missing class docstring (missing-class-docstring) +C0116: Missing function or method docstring (missing-function-docstring) +done: { for f in ./*.py; do [[ "$f" != "./__init__.py" ]] && grep -v '^[[:space:]]*#' "$f"; done; } | pbcopy + +然后是一个完整的tests，但是不要用其他的包，只是test开头的函数或者类，要求from loguru import logger \ No newline at end of file diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py new file mode 100644 index 00000000..bacf342f --- /dev/null +++ b/reme_ai/core/utils/__init__.py @@ -0,0 +1,5 @@ +"""utils""" + +from .timer import timer + +__all__ = ["timer"] diff --git a/reme_ai/core/utils/timer.py b/reme_ai/core/utils/timer.py new file mode 100644 index 00000000..4660ac2f --- /dev/null +++ b/reme_ai/core/utils/timer.py @@ -0,0 +1,66 @@ +""" +Utility module for timing function execution with log metadata preservation. +""" + +import functools +import inspect +import time +from typing import Any, Callable, TypeVar, cast + +from loguru import logger + +# Type variable to preserve the signature of the decorated callable +F = TypeVar("F", bound=Callable[..., Any]) + + +def timer(func: F) -> F: + """ + Decorator that logs execution time and patches log records with original function metadata. + """ + # Extract original function metadata to ensure logs point to the correct source + func_name = func.__name__ + try: + # Retrieve the source file path and the starting line number + file_path = inspect.getsourcefile(func) or "unknown" + _, line_no = inspect.getsourcelines(func) + except Exception: + file_path = "unknown" + line_no = 0 + + def patcher(record): + """Modifies the log record to reflect the decorated function's location.""" + record["function"] = func_name + record["file"].name = file_path.split("/")[-1] + record["file"].path = file_path + record["line"] = line_no + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + """Timer wrapper for asynchronous functions.""" + start_time = time.perf_counter() + try: + return await func(*args, **kwargs) + finally: + duration = time.perf_counter() - start_time + # Use patch to inject metadata instead of relying on stack depth + logger.patch(patcher).info( + "========== cost={:.6f}s ==========", + duration + ) + + @functools.wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: + """Timer wrapper for synchronous functions.""" + start_time = time.perf_counter() + try: + return func(*args, **kwargs) + finally: + duration = time.perf_counter() - start_time + logger.patch(patcher).info( + "========== cost={:.6f}s ==========", + duration + ) + + if inspect.iscoroutinefunction(func): + return cast(F, async_wrapper) + return cast(F, sync_wrapper) diff --git a/tests/test_timer.py b/tests/test_timer.py new file mode 100644 index 00000000..c9714e38 --- /dev/null +++ b/tests/test_timer.py @@ -0,0 +1,64 @@ +""" +This module provides a suite of tests to verify universal timer decorator functionality using loguru. +""" + +import asyncio +import time + +from loguru import logger + +from reme_ai.core.utils import timer + + +@timer +def test_sync_function(seconds: float) -> str: + """Tests timing of a standard synchronous function.""" + time.sleep(seconds) + return "sync done" + + +@timer +async def test_async_function(seconds: float) -> str: + """Tests timing of an asynchronous function.""" + await asyncio.sleep(seconds) + return "async done" + + +class TestMemberMethods: + """Container class to test class method decoration.""" + + @timer + def test_sync_method(self, seconds: float) -> None: + """Tests a synchronous instance method.""" + time.sleep(seconds) + + @timer + async def test_async_method(self, seconds: float) -> None: + """Tests an asynchronous instance method.""" + await asyncio.sleep(seconds) + + +def run_all_tests() -> None: + """ + Manual test runner. + Notice that the logs will now point to the line numbers below + (where the function is actually called). + """ + logger.info("Starting tests and verifying stack trace...") + + # 1. Test Sync Function + test_sync_function(0.1) + + # 2. Test Async Function + asyncio.run(test_async_function(0.1)) + + # 3. Test Class Methods + tester = TestMemberMethods() + tester.test_sync_method(0.05) + asyncio.run(tester.test_async_method(0.05)) + + logger.success("All tests completed.") + + +if __name__ == "__main__": + run_all_tests() From 266b19ecd059c301bc48b00d91bffa6f99e3aef7 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 30 Dec 2025 18:16:33 +0800 Subject: [PATCH 02/10] feat(core): add enumeration and schema modules with utility functions --- docs/deprecated.txt | 3 +- reme_ai/core/enumeration/__init__.py | 13 +++ reme_ai/core/enumeration/chunk_enum.py | 25 +++++ reme_ai/core/enumeration/http_enum.py | 22 ++++ reme_ai/core/enumeration/registry_enum.py | 28 +++++ reme_ai/core/enumeration/role.py | 19 ++++ reme_ai/core/schema/__init__.py | 40 +++++++ reme_ai/core/schema/message.py | 112 +++++++++++++++++++ reme_ai/core/schema/request.py | 17 +++ reme_ai/core/schema/response.py | 11 ++ reme_ai/core/schema/service_config.py | 113 +++++++++++++++++++ reme_ai/core/schema/stream_chunk.py | 14 +++ reme_ai/core/schema/tool_call.py | 128 ++++++++++++++++++++++ reme_ai/core/schema/vector_node.py | 15 +++ reme_ai/core/utils/__init__.py | 3 +- reme_ai/core/utils/case_converter.py | 28 +++++ reme_ai/core/utils/singleton.py | 17 +++ reme_ai/core/utils/timer.py | 4 +- 18 files changed, 608 insertions(+), 4 deletions(-) create mode 100644 reme_ai/core/enumeration/__init__.py create mode 100644 reme_ai/core/enumeration/chunk_enum.py create mode 100644 reme_ai/core/enumeration/http_enum.py create mode 100644 reme_ai/core/enumeration/registry_enum.py create mode 100644 reme_ai/core/enumeration/role.py create mode 100644 reme_ai/core/schema/__init__.py create mode 100644 reme_ai/core/schema/message.py create mode 100644 reme_ai/core/schema/request.py create mode 100644 reme_ai/core/schema/response.py create mode 100644 reme_ai/core/schema/service_config.py create mode 100644 reme_ai/core/schema/stream_chunk.py create mode 100644 reme_ai/core/schema/tool_call.py create mode 100644 reme_ai/core/schema/vector_node.py create mode 100644 reme_ai/core/utils/case_converter.py create mode 100644 reme_ai/core/utils/singleton.py diff --git a/docs/deprecated.txt b/docs/deprecated.txt index 41d9c205..1327a043 100644 --- a/docs/deprecated.txt +++ b/docs/deprecated.txt @@ -1,9 +1,10 @@ from loguru import logger 用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码 +用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范 C0114: Missing module docstring (missing-module-docstring) C0115: Missing class docstring (missing-class-docstring) C0116: Missing function or method docstring (missing-function-docstring) done: { for f in ./*.py; do [[ "$f" != "./__init__.py" ]] && grep -v '^[[:space:]]*#' "$f"; done; } | pbcopy -然后是一个完整的tests，但是不要用其他的包，只是test开头的函数或者类，要求from loguru import logger \ No newline at end of file +然后是一个完整的tests，但是不要用其他的包，只是test开头的函数或者类，要求from loguru import logger diff --git a/reme_ai/core/enumeration/__init__.py b/reme_ai/core/enumeration/__init__.py new file mode 100644 index 00000000..905c2091 --- /dev/null +++ b/reme_ai/core/enumeration/__init__.py @@ -0,0 +1,13 @@ +"""enumeration""" + +from .chunk_enum import ChunkEnum +from .http_enum import HttpEnum +from .registry_enum import RegistryEnum +from .role import Role + +__all__ = [ + "ChunkEnum", + "HttpEnum", + "RegistryEnum", + "Role", +] diff --git a/reme_ai/core/enumeration/chunk_enum.py b/reme_ai/core/enumeration/chunk_enum.py new file mode 100644 index 00000000..dbe37106 --- /dev/null +++ b/reme_ai/core/enumeration/chunk_enum.py @@ -0,0 +1,25 @@ +"""Defines the types of data chunks used in streaming responses.""" + +from enum import Enum + + +class ChunkEnum(str, Enum): + """Enumeration of possible chunk categories for stream processing.""" + + # Internal reasoning or chain-of-thought process + THINK = "think" + + # The final generated response content + ANSWER = "answer" + + # Metadata or calls related to external tools + TOOL = "tool" + + # Resource consumption and token usage statistics + USAGE = "usage" + + # Error messages or exception details + ERROR = "error" + + # Final signal indicating the completion of the stream + DONE = "done" diff --git a/reme_ai/core/enumeration/http_enum.py b/reme_ai/core/enumeration/http_enum.py new file mode 100644 index 00000000..19622242 --- /dev/null +++ b/reme_ai/core/enumeration/http_enum.py @@ -0,0 +1,22 @@ +"""Provides a collection of standard HTTP request methods.""" + +from enum import Enum + + +class HttpEnum(str, Enum): + """Enumeration of supported HTTP methods for network requests.""" + + # Retrieves data from a specified resource + GET = "get" + + # Submits data to be processed to a specified resource + POST = "post" + + # Identical to GET but only retrieves the response headers + HEAD = "head" + + # Uploads or replaces the representation of a target resource + PUT = "put" + + # Deletes the specified resource from the server + DELETE = "delete" diff --git a/reme_ai/core/enumeration/registry_enum.py b/reme_ai/core/enumeration/registry_enum.py new file mode 100644 index 00000000..876c06b8 --- /dev/null +++ b/reme_ai/core/enumeration/registry_enum.py @@ -0,0 +1,28 @@ +"""Defines the registry categories for core components of the system.""" + +from enum import Enum + + +class RegistryEnum(str, Enum): + """Enumeration of component types registered within the application lifecycle.""" + + # Large Language Model interfaces + LLM = "llm" + + # Models used for generating vector embeddings + EMBEDDING_MODEL = "embedding_model" + + # Databases or storage systems for vector search + VECTOR_STORE = "vector_store" + + # Atomic operations or functional units + OP = "op" + + # Orchestrated sequences of operations or workflows + FLOW = "flow" + + # External APIs or shared internal services + SERVICE = "service" + + # Utilities for tracking and limiting token consumption + TOKEN_COUNTER = "token_counter" diff --git a/reme_ai/core/enumeration/role.py b/reme_ai/core/enumeration/role.py new file mode 100644 index 00000000..4acad7e5 --- /dev/null +++ b/reme_ai/core/enumeration/role.py @@ -0,0 +1,19 @@ +"""Defines the participant roles in a chat completion sequence.""" + +from enum import Enum + + +class Role(str, Enum): + """Enumeration of standard personas involved in a conversation flow.""" + + # High-level instructions to guide the model's behavior + SYSTEM = "system" + + # Input or queries provided by the human user + USER = "user" + + # Responses or messages generated by the AI model + ASSISTANT = "assistant" + + # Output or results returned from external tool executions + TOOL = "tool" diff --git a/reme_ai/core/schema/__init__.py b/reme_ai/core/schema/__init__.py new file mode 100644 index 00000000..f55f054e --- /dev/null +++ b/reme_ai/core/schema/__init__.py @@ -0,0 +1,40 @@ +"""schema""" + +from .message import ContentBlock, Message, Trajectory +from .request import Request +from .response import Response +from .service_config import ( + CmdConfig, + EmbeddingModelConfig, + FlowConfig, + HttpConfig, + LLMConfig, + MCPConfig, + ServiceConfig, + TokenCounterConfig, + VectorStoreConfig, +) +from .stream_chunk import StreamChunk +from .tool_call import ToolAttr, ToolCall +from .vector_node import VectorNode + +__all__ = [ + "ContentBlock", + "EmbeddingModelConfig", + "FlowConfig", + "HttpConfig", + "LLMConfig", + "MCPConfig", + "Message", + "Request", + "Response", + "ServiceConfig", + "StreamChunk", + "TokenCounterConfig", + "Trajectory", + "ToolAttr", + "ToolCall", + "VectorNode", + "VectorStoreConfig", + "CmdConfig", +] diff --git a/reme_ai/core/schema/message.py b/reme_ai/core/schema/message.py new file mode 100644 index 00000000..b5c833ea --- /dev/null +++ b/reme_ai/core/schema/message.py @@ -0,0 +1,112 @@ +"""Data models for multi-modal conversation history and LLM interaction trajectories.""" + +import datetime +import json +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from .tool_call import ToolCall +from ..enumeration import Role + + +class ContentBlock(BaseModel): + """Individual unit of multi-modal content like text, images, or video.""" + + model_config = ConfigDict(extra="allow") + + type: str = Field(default="") + content: str | dict | list = Field(default="") + + @model_validator(mode="before") + @classmethod + def init_block(cls, data: dict[str, Any]) -> dict[str, Any]: + """Dynamically maps the type-specific key to the content field.""" + content_type = data.get("type", "") + if content_type and content_type in data: + data["content"] = data[content_type] + return data + + def simple_dump(self) -> dict[str, Any]: + """Serializes the block into an API-compatible dictionary format.""" + return { + "type": self.type, + self.type: self.content, + **self.model_extra, + } + + +class Message(BaseModel): + """Data model for a single dialogue entry including roles and tool interactions.""" + + name: str | None = Field(default=None) + role: Role = Field(default=Role.USER) + content: str | list[ContentBlock] = Field(default="") + reasoning_content: str = Field(default="") + tool_calls: list[ToolCall] = Field(default_factory=list) + tool_call_id: str = Field(default="") + time_created: str = Field( + default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + ) + metadata: dict[str, Any] = Field(default_factory=dict) + + def dump_content(self) -> str | list[dict[str, Any]]: + """Returns content as a raw string or a list of serialized blocks.""" + if isinstance(self.content, str): + return self.content + return [block.simple_dump() for block in self.content] + + def simple_dump(self, add_reasoning: bool = True) -> dict[str, Any]: + """Transforms the message into a simplified dictionary for standard APIs.""" + result = {"role": self.role.value, "content": self.dump_content()} + + if add_reasoning and self.reasoning_content: + result["reasoning_content"] = self.reasoning_content + if self.tool_calls: + result["tool_calls"] = [tc.simple_output_dump() for tc in self.tool_calls] + if self.tool_call_id: + result["tool_call_id"] = self.tool_call_id + + return result + + def format_message( + self, + index: int | None = None, + add_time: bool = False, + use_name: bool = False, + add_reasoning: bool = True, + add_tools: bool = True, + ) -> str: + """Generates a human-readable string representation of the message.""" + prefix = f"round{index} " if index is not None else "" + time_str = f"[{self.time_created}] " if add_time else "" + header = f"{self.name or self.role.value if use_name else self.role.value}:\n" + + lines = [f"{prefix}{time_str}{header}"] + + if add_reasoning and self.reasoning_content: + lines.append(f"{self.reasoning_content}\n") + + if isinstance(self.content, str): + lines.append(self.content) + elif isinstance(self.content, list): + for block in self.content: + text = ( + block.content if isinstance(block.content, str) else json.dumps(block.content, ensure_ascii=False) + ) + lines.append(str(text)) + + if add_tools and self.tool_calls: + for tc in self.tool_calls: + lines.append(f" - tool_call={tc.name} params={tc.arguments}") + + return "\n".join(lines).strip() + + +class Trajectory(BaseModel): + """Sequence of messages representing a full conversation session and its evaluation.""" + + task_id: str = Field(default="") + messages: list[Message] = Field(default_factory=list) + score: float = Field(default=0.0) + metadata: dict[str, Any] = Field(default_factory=dict) diff --git a/reme_ai/core/schema/request.py b/reme_ai/core/schema/request.py new file mode 100644 index 00000000..aa054219 --- /dev/null +++ b/reme_ai/core/schema/request.py @@ -0,0 +1,17 @@ +"""Defines the data structure for processing incoming user requests and message history.""" + +from typing import List + +from pydantic import Field, BaseModel, ConfigDict + +from .message import Message + + +class Request(BaseModel): + """Represents a structured request payload containing a query, message list, and metadata.""" + + model_config = ConfigDict(extra="allow") + + query: str = Field(default="") + messages: List[Message] = Field(default_factory=list) + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/response.py b/reme_ai/core/schema/response.py new file mode 100644 index 00000000..3104bc6e --- /dev/null +++ b/reme_ai/core/schema/response.py @@ -0,0 +1,11 @@ +"""Defines the standardized data structure for model output responses.""" + +from pydantic import Field, BaseModel + + +class Response(BaseModel): + """Represents a structured response containing the execution result, status, and metadata.""" + + answer: str | dict | list = Field(default="") + success: bool = Field(default=True) + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/service_config.py b/reme_ai/core/schema/service_config.py new file mode 100644 index 00000000..67cac68f --- /dev/null +++ b/reme_ai/core/schema/service_config.py @@ -0,0 +1,113 @@ +"""Configuration schemas for service components using Pydantic models.""" + +from typing import Dict, List + +from pydantic import BaseModel, Field, ConfigDict + +from .tool_call import ToolCall + + +class MCPConfig(BaseModel): + """Configuration for Model Context Protocol transport and network settings.""" + + model_config = ConfigDict(extra="allow") + + transport: str = Field(default="") + host: str = Field(default="0.0.0.0") + port: int = Field(default=8001) + + +class HttpConfig(BaseModel): + """Configuration for the HTTP server interface and connection lifecycle.""" + + model_config = ConfigDict(extra="allow") + + host: str = Field(default="0.0.0.0") + port: int = Field(default=8001) + timeout_keep_alive: int = Field(default=3600) + limit_concurrency: int = Field(default=1000) + + +class CmdConfig(BaseModel): + """Configuration for command-line flow execution parameters.""" + + model_config = ConfigDict(extra="allow") + + flow: str = Field(default="") + + +class FlowConfig(ToolCall): + """Configuration for workflow execution, caching, and error handling.""" + + model_config = ConfigDict(extra="allow") + + flow_content: str = Field(default="") + stream: bool = Field(default=False) + raise_exception: bool = Field(default=True) + enable_cache: bool = Field(default=False) + cache_path: str = Field(default="cache/flow") + cache_expire_hours: float = Field(default=0.1) + + +class LLMConfig(BaseModel): + """Configuration for Large Language Model backend and model identification.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + model_name: str = Field(default="") + + +class EmbeddingModelConfig(BaseModel): + """Configuration for embedding model backends and identity.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + model_name: str = Field(default="") + + +class VectorStoreConfig(BaseModel): + """Configuration for vector database storage and associated embeddings.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="local") + collection_name: str = Field(default="remy") + embedding_model: str = Field(default="default") + + +class TokenCounterConfig(BaseModel): + """Configuration for token counting services and model mapping.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="base") + model_name: str = Field(default="") + + +class ServiceConfig(BaseModel): + """Root configuration schema aggregating all service-level settings and components.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="") + enable_logo: bool = Field(default=True) + language: str = Field(default="") + thread_pool_max_workers: int = Field(default=16) + ray_max_workers: int = Field(default=-1) + disabled_flows: List[str] = Field(default_factory=list) + enabled_flows: List[str] = Field(default_factory=list) + external_mcp: Dict[str, dict] = Field( + default_factory=dict, + description="External MCP Server configuration", + ) + + mcp: MCPConfig = Field(default_factory=MCPConfig) + http: HttpConfig = Field(default_factory=HttpConfig) + cmd: CmdConfig = Field(default_factory=CmdConfig) + flow: Dict[str, FlowConfig] = Field(default_factory=dict) + llm: Dict[str, LLMConfig] = Field(default_factory=dict) + embedding_model: Dict[str, EmbeddingModelConfig] = Field(default_factory=dict) + vector_store: Dict[str, VectorStoreConfig] = Field(default_factory=dict) + token_counter: Dict[str, TokenCounterConfig] = Field(default_factory=dict) diff --git a/reme_ai/core/schema/stream_chunk.py b/reme_ai/core/schema/stream_chunk.py new file mode 100644 index 00000000..764981fd --- /dev/null +++ b/reme_ai/core/schema/stream_chunk.py @@ -0,0 +1,14 @@ +"""Defines the data structure for individual data packets in a streaming response.""" + +from pydantic import Field, BaseModel + +from ..enumeration import ChunkEnum + + +class StreamChunk(BaseModel): + """Represents a single chunk of streamed data including its type, content, and completion status.""" + + chunk_type: ChunkEnum = Field(default=ChunkEnum.ANSWER) + chunk: str | dict | list = Field(default="") + done: bool = Field(default=False) + metadata: dict = Field(default_factory=dict) diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py new file mode 100644 index 00000000..a0f77281 --- /dev/null +++ b/reme_ai/core/schema/tool_call.py @@ -0,0 +1,128 @@ +"""Model definitions for MCP tools and LLM tool call interactions.""" + +import json +from typing import Dict, List, Literal, Optional, Any + +from mcp.types import Tool +from pydantic import BaseModel, Field, model_validator, ConfigDict + +TOOL_ATTR_TYPE = Literal["string", "array", "integer", "number", "boolean", "object"] + + +class ToolAttr(BaseModel): + """Represent attributes for tool parameters in a JSON schema format.""" + + type: TOOL_ATTR_TYPE = Field(default="string", description="Attribute data type") + description: str = Field(default="", description="Attribute purpose") + required: bool = Field(default=True, description="Whether the attribute is mandatory") + enum: Optional[List[str]] = Field(default=None, description="Allowed values") + items: Dict[str, Any] = Field(default_factory=dict, description="Schema for array items") + + model_config = ConfigDict(extra="allow") + + def simple_input_dump(self) -> dict: + """Export attribute as a standard JSON schema property dictionary.""" + res: dict = {"type": self.type, "description": self.description} + if self.enum: + res["enum"] = self.enum + if self.items: + res["items"] = self.items + return res + + +class ToolCall(BaseModel): + """Handle tool definitions and execution arguments for LLM integrations.""" + + index: int = Field(default=0) + id: str = Field(default="") + type: str = Field(default="function") + name: str = Field(default="") + arguments: str = Field(default="{}", description="JSON string of execution arguments") + description: str = Field(default="") + input_schema: Dict[str, ToolAttr] = Field(default_factory=dict) + output_schema: Dict[str, ToolAttr] = Field(default_factory=dict) + + @model_validator(mode="before") + @classmethod + def init_tool_call(cls, data: Dict[str, Any]) -> Dict[str, Any]: + """Map raw API response data to the internal ToolCall structure.""" + res = data.copy() + t_type = res.get("type", "function") + inner = res.get(t_type, {}) + + # Extract basic function metadata + for key in ("name", "arguments", "description"): + if key in inner: + res[key] = inner[key] + + # Parse JSON schema parameters into ToolAttr objects + params = inner.get("parameters", {}) + if params: + props = params.get("properties", {}) + reqs = params.get("required", []) + res["input_schema"] = {k: ToolAttr(**v, required=k in reqs) for k, v in props.items()} + return res + + @property + def argument_dict(self) -> dict: + """Parse the arguments string into a dictionary.""" + return json.loads(self.arguments) + + def check_argument(self) -> bool: + """Verify if the arguments string is valid JSON.""" + try: + _ = self.argument_dict + return True + except (json.JSONDecodeError, TypeError): + return False + + @staticmethod + def _build_schema_dict(schema: Dict[str, ToolAttr]) -> dict: + """Construct a JSON schema object from a dictionary of ToolAttrs.""" + return { + "type": "object", + "properties": {k: v.simple_input_dump() for k, v in schema.items()}, + "required": [k for k, v in schema.items() if v.required], + } + + def simple_input_dump(self) -> dict: + """Format the tool definition for LLM provider API requests.""" + return { + "type": self.type, + self.type: { + "name": self.name, + "description": self.description, + "parameters": self._build_schema_dict(self.input_schema), + }, + } + + def simple_output_dump(self) -> dict: + """Format the tool call result for LLM provider API responses.""" + return { + "index": self.index, + "id": self.id, + "type": self.type, + self.type: {"arguments": self.arguments, "name": self.name}, + } + + @classmethod + def from_mcp_tool(cls, tool: Tool) -> "ToolCall": + """Create a ToolCall instance from an MCP Tool object.""" + props = tool.inputSchema.get("properties", {}) + reqs = tool.inputSchema.get("required", []) + return cls( + name=tool.name, + description=tool.description or "", + input_schema={k: ToolAttr(**v, required=k in reqs) for k, v in props.items()}, + ) + + def to_mcp_tool(self) -> Tool: + """Convert the current instance into an MCP Tool object.""" + kwargs = { + "name": self.name, + "description": self.description, + "inputSchema": self._build_schema_dict(self.input_schema), + } + if self.output_schema: + kwargs["outputSchema"] = self._build_schema_dict(self.output_schema) + return Tool(**kwargs) diff --git a/reme_ai/core/schema/vector_node.py b/reme_ai/core/schema/vector_node.py new file mode 100644 index 00000000..937ef4be --- /dev/null +++ b/reme_ai/core/schema/vector_node.py @@ -0,0 +1,15 @@ +"""Defines the data structure for individual vector embedding nodes within a retrieval system.""" + +from typing import List, Dict +from uuid import uuid4 + +from pydantic import BaseModel, Field + + +class VectorNode(BaseModel): + """Represents a discrete unit of text content paired with its corresponding vector embedding and metadata.""" + + vector_id: str = Field(default_factory=lambda: uuid4().hex) + content: str = Field(default="") + vector: List[float] | None = Field(default=None) + metadata: Dict[str, str | bool | int | float] = Field(default_factory=dict) diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index bacf342f..62a9d5a0 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -1,5 +1,6 @@ """utils""" from .timer import timer +from .singleton import singleton -__all__ = ["timer"] +__all__ = ["timer", "singleton"] diff --git a/reme_ai/core/utils/case_converter.py b/reme_ai/core/utils/case_converter.py new file mode 100644 index 00000000..28815282 --- /dev/null +++ b/reme_ai/core/utils/case_converter.py @@ -0,0 +1,28 @@ +"""Case conversion utility for PascalCase, camelCase, and snake_case.""" + +import re + +# Acronyms that should remain uppercase in Pascal/camelCase +_ACRONYMS = {"LLM", "API", "URL", "HTTP", "JSON", "XML", "AI", "MCP"} +_ACRONYM_MAP = {word.lower(): word for word in _ACRONYMS} + + +def camel_to_snake(content: str) -> str: + """Convert PascalCase or camelCase to snake_case.""" + # Normalize acronyms to title case (e.g., LLM -> Llm) to assist regex splitting + for word in _ACRONYMS: + content = content.replace(word, word.capitalize()) + + # Insert underscores between case transitions and convert to lowercase + return re.sub(r"(? str: + """Convert snake_case to PascalCase (preserving defined acronyms).""" + return "".join(_ACRONYM_MAP.get(part.lower(), part.capitalize()) for part in content.split("_") if part) + + +if __name__ == "__main__": + # Quick verification + print(camel_to_snake("OpenAILLMClient")) # open_ai_llm_client + print(snake_to_camel("open_ai_llm_client")) # OpenAILLMClient diff --git a/reme_ai/core/utils/singleton.py b/reme_ai/core/utils/singleton.py new file mode 100644 index 00000000..0f6c5907 --- /dev/null +++ b/reme_ai/core/utils/singleton.py @@ -0,0 +1,17 @@ +"""Module providing a decorator to implement the Singleton design pattern.""" + + +def singleton(cls): + """A class decorator that ensures only one instance of a class exists.""" + + # Dictionary to cache the single instance of the class + _instance = {} + + def _singleton(*args, **kwargs): + """Return the existing instance or create a new one if it doesn't exist.""" + if cls not in _instance: + # Create and store the instance if it's the first call + _instance[cls] = cls(*args, **kwargs) + return _instance[cls] + + return _singleton diff --git a/reme_ai/core/utils/timer.py b/reme_ai/core/utils/timer.py index 4660ac2f..f03225fa 100644 --- a/reme_ai/core/utils/timer.py +++ b/reme_ai/core/utils/timer.py @@ -45,7 +45,7 @@ def timer(func: F) -> F: # Use patch to inject metadata instead of relying on stack depth logger.patch(patcher).info( "========== cost={:.6f}s ==========", - duration + duration, ) @functools.wraps(func) @@ -58,7 +58,7 @@ def timer(func: F) -> F: duration = time.perf_counter() - start_time logger.patch(patcher).info( "========== cost={:.6f}s ==========", - duration + duration, ) if inspect.iscoroutinefunction(func): From e5f17d6e67e8d7905184fc08269074866a367716 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 30 Dec 2025 22:01:56 +0800 Subject: [PATCH 03/10] docs(guidelines): add test file creation guidelines with loguru logger --- docs/deprecated.txt | 1 + reme_ai/core/schema/message.py | 26 ++++- reme_ai/core/schema/tool_call.py | 38 +++++++- tests/test_message.py | 161 +++++++++++++++++++++++++++++++ 4 files changed, 221 insertions(+), 5 deletions(-) create mode 100644 tests/test_message.py diff --git a/docs/deprecated.txt b/docs/deprecated.txt index 1327a043..6d7cb14f 100644 --- a/docs/deprecated.txt +++ b/docs/deprecated.txt @@ -8,3 +8,4 @@ C0116: Missing function or method docstring (missing-function-docstring) done: { for f in ./*.py; do [[ "$f" != "./__init__.py" ]] && grep -v '^[[:space:]]*#' "$f"; done; } | pbcopy 然后是一个完整的tests，但是不要用其他的包，只是test开头的函数或者类，要求from loguru import logger +写一个测试文件，不要使用pytest，普通的test，要求英文注释 \ No newline at end of file diff --git a/reme_ai/core/schema/message.py b/reme_ai/core/schema/message.py index b5c833ea..496ef7ec 100644 --- a/reme_ai/core/schema/message.py +++ b/reme_ai/core/schema/message.py @@ -11,7 +11,31 @@ from ..enumeration import Role class ContentBlock(BaseModel): - """Individual unit of multi-modal content like text, images, or video.""" + """ + Individual unit of multi-modal content like text, images, or video. + examples: + { + "type": "image_url", + "image_url": { + "url": "https://img.alicdn.com/imgextra/i1/O1CN01gDEY8M1W114Hi3XcN_!!6000000002727-0-tps-1024-406.jpg" + }, + } + + { + "type": "video", + "video": [ + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/xzsgiz/football1.jpg", + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/tdescd/football2.jpg", + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/zefdja/football3.jpg", + "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/aedbqh/football4.jpg", + ], + } + + { + "type": "text", + "text": "How do you solve this problem?" + } + """ model_config = ConfigDict(extra="allow") diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index a0f77281..c0d544df 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core/schema/tool_call.py @@ -1,7 +1,7 @@ """Model definitions for MCP tools and LLM tool call interactions.""" import json -from typing import Dict, List, Literal, Optional, Any +from typing import Dict, Literal, Any from mcp.types import Tool from pydantic import BaseModel, Field, model_validator, ConfigDict @@ -15,8 +15,8 @@ class ToolAttr(BaseModel): type: TOOL_ATTR_TYPE = Field(default="string", description="Attribute data type") description: str = Field(default="", description="Attribute purpose") required: bool = Field(default=True, description="Whether the attribute is mandatory") - enum: Optional[List[str]] = Field(default=None, description="Allowed values") - items: Dict[str, Any] = Field(default_factory=dict, description="Schema for array items") + enum: list[str] | None = Field(default=None, description="Allowed values") + items: dict[str, Any] | None = Field(default=None, description="Schema for array items") model_config = ConfigDict(extra="allow") @@ -31,7 +31,37 @@ class ToolAttr(BaseModel): class ToolCall(BaseModel): - """Handle tool definitions and execution arguments for LLM integrations.""" + """ + Handle tool definitions and execution arguments for LLM integrations. + input: + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "It is very useful when you want to check the weather of a specified city.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.", + } + }, + "required": ["location"] + } + } + } + output: + { + "index": 0, + "id": "call_6596dafa2a6a46f7a217da", + "function": { + "arguments": "{\"location\": \"Beijing\"}", + "name": "get_current_weather" + }, + "type": "function", + } + """ index: int = Field(default=0) id: str = Field(default="") diff --git a/tests/test_message.py b/tests/test_message.py new file mode 100644 index 00000000..14f4c52f --- /dev/null +++ b/tests/test_message.py @@ -0,0 +1,161 @@ +"""Test cases for message schema and serialization.""" + +import unittest + +from mcp.types import Tool + +from reme_ai.core.enumeration import Role +from reme_ai.core.schema import ToolAttr, ToolCall, ContentBlock, Message + + +class TestModelDefinitions(unittest.TestCase): + """Test suite for validating message models and their serialization methods.""" + + def test_tool_attr_serialization(self): + """Test if ToolAttr correctly dumps to JSON schema format.""" + attr = ToolAttr( + type="string", + description="The city name", + enum=["Beijing", "London"], + required=True, + ) + dump = attr.simple_input_dump() + + print("\n=== ToolAttr.simple_input_dump() ===") + print(dump) + + self.assertEqual(dump["type"], "string") + self.assertEqual(dump["enum"], ["Beijing", "London"]) + self.assertIn("description", dump) + + def test_tool_call_initialization(self): + """Test if ToolCall correctly parses raw OpenAI-style tool definitions.""" + raw_input = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Check weather info", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"}, + }, + "required": ["location"], + }, + }, + } + + tc = ToolCall(**raw_input) + + print("\n=== ToolCall.simple_input_dump() ===") + print(tc.simple_input_dump()) + print("\n=== ToolCall.simple_output_dump() ===") + print(tc.simple_output_dump()) + + self.assertEqual(tc.name, "get_weather") + self.assertIn("location", tc.input_schema) + self.assertTrue(tc.input_schema["location"].required) + + def test_tool_call_argument_parsing(self): + """Test JSON argument parsing and validation.""" + tc = ToolCall(name="test", arguments='{"key": "value"}') + + self.assertTrue(tc.check_argument()) + self.assertEqual(tc.argument_dict["key"], "value") + + # Test invalid JSON + tc.arguments = "{invalid_json}" + self.assertFalse(tc.check_argument()) + + def test_content_block_dynamic_mapping(self): + """Test if ContentBlock correctly identifies content based on type key.""" + # Test Image Block + img_data = {"type": "image_url", "image_url": {"url": "http://test.com/a.jpg"}} + block = ContentBlock(**img_data) + self.assertEqual(block.type, "image_url") + self.assertEqual(block.content["url"], "http://test.com/a.jpg") + + # Test Text Block + text_data = {"type": "text", "text": "Hello World"} + block = ContentBlock(**text_data) + self.assertEqual(block.content, "Hello World") + + def test_message_simple_dump(self): + """Test the transformation of Message to standard API dictionary.""" + msg = Message( + role=Role.ASSISTANT, + content="Thinking...", + reasoning_content="I should check the weather first.", + tool_calls=[ToolCall(name="get_weather", arguments='{"city": "NY"}', id="call_123")], + ) + + dump = msg.simple_dump(add_reasoning=True) + + print("\n=== Message.simple_dump(add_reasoning=True) ===") + print(dump) + + dump_no_reasoning = msg.simple_dump(add_reasoning=False) + print("\n=== Message.simple_dump(add_reasoning=False) ===") + print(dump_no_reasoning) + + self.assertEqual(dump["role"], "assistant") + self.assertEqual(dump["reasoning_content"], "I should check the weather first.") + self.assertEqual(len(dump["tool_calls"]), 1) + self.assertEqual(dump["tool_calls"][0]["id"], "call_123") + + def test_message_format_human_readable(self): + """Test the string representation of messages for logging/UI.""" + msg = Message( + role=Role.USER, + content=[ + ContentBlock(type="text", text="Look at this:"), + ContentBlock(type="image_url", image_url={"url": "img.png"}), + ], + ) + + formatted = msg.format_message(index=1, use_name=False) + + self.assertIn("round1", formatted) + self.assertIn("user:", formatted) + self.assertIn("Look at this:", formatted) + self.assertIn("img.png", formatted) + + def test_mcp_conversion(self): + """Test the interoperability with MCP Tool format.""" + # Create a mock MCP Tool + mcp_tool = Tool( + name="calculator", + description="adds numbers", + inputSchema={ + "type": "object", + "properties": {"a": {"type": "number"}}, + "required": ["a"], + }, + ) + + # From MCP to ToolCall + tc = ToolCall.from_mcp_tool(mcp_tool) + + print("\n=== ToolCall from MCP - simple_input_dump() ===") + print(tc.simple_input_dump()) + print("\n=== ToolCall from MCP - simple_output_dump() ===") + print(tc.simple_output_dump()) + + self.assertEqual(tc.name, "calculator") + self.assertTrue(tc.input_schema["a"].required) + + # From ToolCall back to MCP structure (via to_mcp_tool) + # Note: This checks the logic of constructing the dict for Tool(...) + mcp_compatible = tc.to_mcp_tool() + + print("\n=== MCP Tool converted back ===") + print(f"Name: {mcp_compatible.name}") + print(f"Description: {mcp_compatible.description}") + print(f"InputSchema: {mcp_compatible.inputSchema}") + + self.assertEqual(mcp_compatible.name, "calculator") + self.assertIn("a", mcp_compatible.inputSchema["properties"]) + + +if __name__ == "__main__": + unittest.main() From ad6395f01278d6826b885db3559c6a354dc853a3 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 30 Dec 2025 23:52:45 +0800 Subject: [PATCH 04/10] feat(schema): add JSON Schema enum and enhance ToolCall with recursive schema support --- reme_ai/core/enumeration/__init__.py | 2 + reme_ai/core/enumeration/json_schema_enum.py | 19 + reme_ai/core/schema/tool_call.py | 217 ++++++----- tests/test_message.py | 44 ++- tests/test_tool_call.py | 358 +++++++++++++++++++ 5 files changed, 524 insertions(+), 116 deletions(-) create mode 100644 reme_ai/core/enumeration/json_schema_enum.py create mode 100644 tests/test_tool_call.py diff --git a/reme_ai/core/enumeration/__init__.py b/reme_ai/core/enumeration/__init__.py index 905c2091..3323bdc9 100644 --- a/reme_ai/core/enumeration/__init__.py +++ b/reme_ai/core/enumeration/__init__.py @@ -2,12 +2,14 @@ from .chunk_enum import ChunkEnum from .http_enum import HttpEnum +from .json_schema_enum import JsonSchemaEnum from .registry_enum import RegistryEnum from .role import Role __all__ = [ "ChunkEnum", "HttpEnum", + "JsonSchemaEnum", "RegistryEnum", "Role", ] diff --git a/reme_ai/core/enumeration/json_schema_enum.py b/reme_ai/core/enumeration/json_schema_enum.py new file mode 100644 index 00000000..17b59380 --- /dev/null +++ b/reme_ai/core/enumeration/json_schema_enum.py @@ -0,0 +1,19 @@ +"""Defines the standard data types supported by JSON Schema.""" + +from enum import Enum + + +class JsonSchemaEnum(str, Enum): + """Enumeration of valid JSON Schema data types.""" + + STRING = "string" + NUMBER = "number" + INTEGER = "integer" + OBJECT = "object" + ARRAY = "array" + BOOLEAN = "boolean" + NULL = "null" + + def __str__(self) -> str: + """Returns the string representation of the enum value.""" + return self.value diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index c0d544df..3688da54 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core/schema/tool_call.py @@ -1,158 +1,153 @@ -"""Model definitions for MCP tools and LLM tool call interactions.""" - +""" +MCP Tool Schema definitions for recursive JSON Schema representation. +""" import json -from typing import Dict, Literal, Any +from typing import Any, Dict, List, Literal, Optional, Union from mcp.types import Tool -from pydantic import BaseModel, Field, model_validator, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, model_validator -TOOL_ATTR_TYPE = Literal["string", "array", "integer", "number", "boolean", "object"] +from ..enumeration.json_schema_enum import JsonSchemaEnum class ToolAttr(BaseModel): - """Represent attributes for tool parameters in a JSON schema format.""" - - type: TOOL_ATTR_TYPE = Field(default="string", description="Attribute data type") - description: str = Field(default="", description="Attribute purpose") - required: bool = Field(default=True, description="Whether the attribute is mandatory") - enum: list[str] | None = Field(default=None, description="Allowed values") - items: dict[str, Any] | None = Field(default=None, description="Schema for array items") - + """Recursive model representing JSON Schema attributes for tool parameters.""" model_config = ConfigDict(extra="allow") + type: Literal[ + JsonSchemaEnum.STRING.value, + JsonSchemaEnum.NUMBER.value, + JsonSchemaEnum.INTEGER.value, + JsonSchemaEnum.OBJECT.value, + JsonSchemaEnum.ARRAY.value, + JsonSchemaEnum.BOOLEAN.value, + JsonSchemaEnum.NULL.value, + ] = Field( + default=JsonSchemaEnum.STRING.value, + description="The data type of the attribute" + ) + description: Optional[str] = Field(default=None, description="Description of the attribute") + required: Optional[List[str]] = Field(default=None, description="Required property names for object types") + properties: Optional[Dict[str, "ToolAttr"]] = Field(default=None, description="Child properties for objects") + items: Optional[Union[Dict[str, Any], "ToolAttr"]] = Field(default=None, description="Schema for array items") + enum: Optional[List[str]] = Field(default=None, description="Allowed values for the attribute") + + def simple_input_dump(self) -> dict: - """Export attribute as a standard JSON schema property dictionary.""" - res: dict = {"type": self.type, "description": self.description} + """Serializes the attribute into a standard JSON Schema dictionary.""" + res: dict = {"type": self.type} + if self.description: + res["description"] = self.description if self.enum: res["enum"] = self.enum - if self.items: - res["items"] = self.items + + if self.type == "object" and self.properties: + res["properties"] = {k: v.simple_input_dump() if isinstance(v, ToolAttr) else v + for k, v in self.properties.items()} + if self.required: + res["required"] = self.required + + if self.type == "array" and self.items: + res["items"] = self.items.simple_input_dump() if isinstance(self.items, ToolAttr) else self.items + return res -class ToolCall(BaseModel): - """ - Handle tool definitions and execution arguments for LLM integrations. - input: - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "It is very useful when you want to check the weather of a specified city.", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.", - } - }, - "required": ["location"] - } - } - } - output: - { - "index": 0, - "id": "call_6596dafa2a6a46f7a217da", - "function": { - "arguments": "{\"location\": \"Beijing\"}", - "name": "get_current_weather" - }, - "type": "function", - } - """ +# Enable recursive type resolution +ToolAttr.model_rebuild() - index: int = Field(default=0) - id: str = Field(default="") - type: str = Field(default="function") - name: str = Field(default="") - arguments: str = Field(default="{}", description="JSON string of execution arguments") - description: str = Field(default="") + +class ToolCall(BaseModel): + """Model representing a tool definition and its call structure.""" + + index: int = 0 + id: str = "" + type: str = "function" + name: str = "" + arguments: str = Field(default="", description="JSON string of tool execution arguments") + description: str = "" input_schema: Dict[str, ToolAttr] = Field(default_factory=dict) + input_required: List[str] = Field(default_factory=list) output_schema: Dict[str, ToolAttr] = Field(default_factory=dict) @model_validator(mode="before") @classmethod - def init_tool_call(cls, data: Dict[str, Any]) -> Dict[str, Any]: - """Map raw API response data to the internal ToolCall structure.""" - res = data.copy() - t_type = res.get("type", "function") - inner = res.get(t_type, {}) + def init_tool_call(cls, data: dict) -> dict: + """Initializes the model by parsing tool-specific body data.""" + data = data.copy() + t_type = data.get("type", "function") + body = data.get(t_type, {}) - # Extract basic function metadata - for key in ("name", "arguments", "description"): - if key in inner: - res[key] = inner[key] + data["name"] = body.get("name", data.get("name", "")) + data["arguments"] = body.get("arguments", data.get("arguments", "")) + data["description"] = body.get("description", data.get("description", "")) - # Parse JSON schema parameters into ToolAttr objects - params = inner.get("parameters", {}) - if params: - props = params.get("properties", {}) - reqs = params.get("required", []) - res["input_schema"] = {k: ToolAttr(**v, required=k in reqs) for k, v in props.items()} - return res + if "parameters" in body: + params = body["parameters"] + data["input_required"] = params.get("required", []) + data["input_schema"] = {k: ToolAttr(**v) for k, v in params.get("properties", {}).items()} - @property - def argument_dict(self) -> dict: - """Parse the arguments string into a dictionary.""" - return json.loads(self.arguments) + return data - def check_argument(self) -> bool: - """Verify if the arguments string is valid JSON.""" - try: - _ = self.argument_dict - return True - except (json.JSONDecodeError, TypeError): - return False - - @staticmethod - def _build_schema_dict(schema: Dict[str, ToolAttr]) -> dict: - """Construct a JSON schema object from a dictionary of ToolAttrs.""" + def _build_full_schema(self) -> dict: + """Generates the top-level JSON Schema object for tool parameters.""" return { "type": "object", - "properties": {k: v.simple_input_dump() for k, v in schema.items()}, - "required": [k for k, v in schema.items() if v.required], + "properties": {k: v.simple_input_dump() for k, v in self.input_schema.items()}, + "required": self.input_required, } def simple_input_dump(self) -> dict: - """Format the tool definition for LLM provider API requests.""" + """Returns a standardized tool definition dictionary.""" return { "type": self.type, self.type: { "name": self.name, "description": self.description, - "parameters": self._build_schema_dict(self.input_schema), + "parameters": self._build_full_schema(), }, } - def simple_output_dump(self) -> dict: - """Format the tool call result for LLM provider API responses.""" - return { - "index": self.index, - "id": self.id, - "type": self.type, - self.type: {"arguments": self.arguments, "name": self.name}, - } - @classmethod def from_mcp_tool(cls, tool: Tool) -> "ToolCall": - """Create a ToolCall instance from an MCP Tool object.""" - props = tool.inputSchema.get("properties", {}) - reqs = tool.inputSchema.get("required", []) + """Creates a ToolCall instance from an MCP Tool object.""" + schema = tool.inputSchema return cls( name=tool.name, description=tool.description or "", - input_schema={k: ToolAttr(**v, required=k in reqs) for k, v in props.items()}, + input_schema={k: ToolAttr(**v) for k, v in schema.get("properties", {}).items()}, + input_required=schema.get("required", []), ) def to_mcp_tool(self) -> Tool: - """Convert the current instance into an MCP Tool object.""" - kwargs = { - "name": self.name, - "description": self.description, - "inputSchema": self._build_schema_dict(self.input_schema), + """Converts the instance back into an MCP Tool object.""" + return Tool( + name=self.name, + description=self.description, + inputSchema=self._build_full_schema(), + ) + + @property + def argument_dict(self) -> dict: + """Parse and return arguments as a dictionary.""" + return json.loads(self.arguments) + + def check_argument(self) -> bool: + """Check if arguments can be parsed as valid JSON.""" + try: + _ = self.argument_dict + return True + except Exception: + return False + + def simple_output_dump(self) -> dict: + """Convert ToolCall to output format dictionary for API responses.""" + return { + "index": self.index, + "id": self.id, + self.type: { + "arguments": self.arguments, + "name": self.name, + }, + "type": self.type, } - if self.output_schema: - kwargs["outputSchema"] = self._build_schema_dict(self.output_schema) - return Tool(**kwargs) diff --git a/tests/test_message.py b/tests/test_message.py index 14f4c52f..35d8598e 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -13,20 +13,41 @@ class TestModelDefinitions(unittest.TestCase): def test_tool_attr_serialization(self): """Test if ToolAttr correctly dumps to JSON schema format.""" + # Test simple string attribute with enum attr = ToolAttr( type="string", description="The city name", enum=["Beijing", "London"], - required=True, ) dump = attr.simple_input_dump() - print("\n=== ToolAttr.simple_input_dump() ===") + print("\n=== ToolAttr.simple_input_dump() (string with enum) ===") print(dump) self.assertEqual(dump["type"], "string") self.assertEqual(dump["enum"], ["Beijing", "London"]) self.assertIn("description", dump) + + # Test object attribute with required child properties + obj_attr = ToolAttr( + type="object", + description="User information", + properties={ + "name": ToolAttr(type="string", description="User name"), + "age": ToolAttr(type="number", description="User age"), + }, + required=["name"], # 'name' is required, 'age' is optional + ) + obj_dump = obj_attr.simple_input_dump() + + print("\n=== ToolAttr.simple_input_dump() (object with required) ===") + print(obj_dump) + + self.assertEqual(obj_dump["type"], "object") + self.assertIn("properties", obj_dump) + self.assertEqual(obj_dump["required"], ["name"]) + self.assertIn("name", obj_dump["properties"]) + self.assertIn("age", obj_dump["properties"]) def test_tool_call_initialization(self): """Test if ToolCall correctly parses raw OpenAI-style tool definitions.""" @@ -39,6 +60,7 @@ class TestModelDefinitions(unittest.TestCase): "type": "object", "properties": { "location": {"type": "string", "description": "City name"}, + "unit": {"type": "string", "description": "Temperature unit"}, }, "required": ["location"], }, @@ -54,7 +76,10 @@ class TestModelDefinitions(unittest.TestCase): self.assertEqual(tc.name, "get_weather") self.assertIn("location", tc.input_schema) - self.assertTrue(tc.input_schema["location"].required) + self.assertIn("unit", tc.input_schema) + # Check that 'location' is in the required list at ToolCall level + self.assertIn("location", tc.input_required) + self.assertNotIn("unit", tc.input_required) def test_tool_call_argument_parsing(self): """Test JSON argument parsing and validation.""" @@ -128,7 +153,10 @@ class TestModelDefinitions(unittest.TestCase): description="adds numbers", inputSchema={ "type": "object", - "properties": {"a": {"type": "number"}}, + "properties": { + "a": {"type": "number", "description": "First number"}, + "b": {"type": "number", "description": "Second number"}, + }, "required": ["a"], }, ) @@ -142,7 +170,11 @@ class TestModelDefinitions(unittest.TestCase): print(tc.simple_output_dump()) self.assertEqual(tc.name, "calculator") - self.assertTrue(tc.input_schema["a"].required) + self.assertIn("a", tc.input_schema) + self.assertIn("b", tc.input_schema) + # Check that 'a' is in the required list + self.assertIn("a", tc.input_required) + self.assertNotIn("b", tc.input_required) # From ToolCall back to MCP structure (via to_mcp_tool) # Note: This checks the logic of constructing the dict for Tool(...) @@ -155,6 +187,8 @@ class TestModelDefinitions(unittest.TestCase): self.assertEqual(mcp_compatible.name, "calculator") self.assertIn("a", mcp_compatible.inputSchema["properties"]) + self.assertIn("b", mcp_compatible.inputSchema["properties"]) + self.assertEqual(mcp_compatible.inputSchema["required"], ["a"]) if __name__ == "__main__": diff --git a/tests/test_tool_call.py b/tests/test_tool_call.py new file mode 100644 index 00000000..242f07ee --- /dev/null +++ b/tests/test_tool_call.py @@ -0,0 +1,358 @@ +import json + +from reme_ai.core.schema.tool_call import ToolCall + + +def test_simple_schema(): + """测试简单的工具定义：只有基本类型参数""" + print("\n========== 测试简单 Schema ==========") + + raw_definition = { + "type": "function", + "function": { + "name": "get_weather", + "description": "获取指定城市的天气信息", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "城市名称"}, + "unit": {"type": "string", "description": "温度单位", "enum": ["celsius", "fahrenheit"]} + }, + "required": ["city"] + } + } + } + + # 解析 + tool_call = ToolCall.model_validate(raw_definition) + print(f"工具名称: {tool_call.name}") + print(f"必填参数: {tool_call.input_required}") + + # 导出并验证相等性 + dumped_data = tool_call.simple_input_dump() + print(f"\n原始定义:\n{json.dumps(raw_definition, indent=2, ensure_ascii=False)}") + print(f"\n导出结果:\n{json.dumps(dumped_data, indent=2, ensure_ascii=False)}") + + # 验证相等 + assert dumped_data == raw_definition, "简单 Schema 导出结果与原始定义不一致" + print("\n✅ 简单 Schema 测试通过：raw_definition == simple_input_dump()") + + +def test_medium_nested_schema(): + """测试中等复杂度：包含一层对象嵌套""" + print("\n========== 测试中等复杂 Schema ==========") + + raw_definition = { + "type": "function", + "function": { + "name": "create_order", + "description": "创建订单", + "parameters": { + "type": "object", + "properties": { + "order_id": {"type": "string", "description": "订单ID"}, + "amount": {"type": "number", "description": "订单金额"}, + "customer": { + "type": "object", + "description": "客户信息", + "properties": { + "name": {"type": "string", "description": "客户姓名"}, + "email": {"type": "string", "description": "客户邮箱"}, + "phone": {"type": "string", "description": "联系电话"} + }, + "required": ["name", "email"] + } + }, + "required": ["order_id", "customer"] + } + } + } + + # 解析 + tool_call = ToolCall.model_validate(raw_definition) + print(f"工具名称: {tool_call.name}") + print(f"根级必填项: {tool_call.input_required}") + + customer_attr = tool_call.input_schema["customer"] + print(f"Customer 子属性: {list(customer_attr.properties.keys())}") + print(f"Customer 必填项: {customer_attr.required}") + + # 导出并验证相等性 + dumped_data = tool_call.simple_input_dump() + print(f"\n原始定义:\n{json.dumps(raw_definition, indent=2, ensure_ascii=False)}") + print(f"\n导出结果:\n{json.dumps(dumped_data, indent=2, ensure_ascii=False)}") + + # 验证相等 + assert dumped_data == raw_definition, "中等复杂 Schema 导出结果与原始定义不一致" + print("\n✅ 中等复杂 Schema 测试通过：raw_definition == simple_input_dump()") + + +def test_nested_schema(): + """测试复杂嵌套：包含对象嵌套和数组嵌套""" + print("\n========== 测试复杂嵌套 Schema ==========") + + # 1. 模拟一个来自 LLM 或 MCP 的复杂嵌套定义 + raw_definition = { + "type": "function", + "function": { + "name": "register_user", + "description": "注册新用户，包含复杂的元数据和标签", + "parameters": { + "type": "object", + "properties": { + "username": {"type": "string", "description": "用户名"}, + "metadata": { + "type": "object", + "description": "用户元数据", + "properties": { + "age": {"type": "integer"}, + "location": {"type": "string"} + }, + "required": ["age"] + }, + "tags": { + "type": "array", + "description": "用户标签列表", + "items": { + "type": "object", + "properties": { + "tag_id": {"type": "string"}, + "level": {"type": "number"} + }, + "required": ["tag_id"] + } + } + }, + "required": ["username", "metadata"] + } + } + } + + # 2. 解析：将原始字典转化为 ToolCall 实例 + tool_call = ToolCall.model_validate(raw_definition) + + print(f"工具名称: {tool_call.name}") + print(f"根级必填项: {tool_call.input_required}") + + # 验证嵌套深度 + metadata_attr = tool_call.input_schema["metadata"] + print(f"Metadata 子属性: {list(metadata_attr.properties.keys())}") + print(f"Metadata 必填项: {metadata_attr.required}") + + # 3. 导出：验证 simple_input_dump 是否生成了正确的 JSON Schema + dumped_data = tool_call.simple_input_dump() + + print(f"\n原始定义:\n{json.dumps(raw_definition, indent=2, ensure_ascii=False)}") + print(f"\n导出结果:\n{json.dumps(dumped_data, indent=2, ensure_ascii=False)}") + + # 验证相等 + assert dumped_data == raw_definition, "复杂嵌套 Schema 导出结果与原始定义不一致" + print("\n✅ 复杂嵌套 Schema 测试通过：raw_definition == simple_input_dump()") + + # 4. 转换验证：测试 to_mcp_tool + mcp_tool = tool_call.to_mcp_tool() + assert mcp_tool.name == "register_user" + assert "properties" in mcp_tool.inputSchema["properties"]["tags"]["items"] + print("✅ 嵌套结构在 MCP Tool 转换中成功保留") + + +def test_array_of_primitives(): + """测试数组嵌套：数组元素为基本类型""" + print("\n========== 测试基本类型数组 Schema ==========") + + raw_definition = { + "type": "function", + "function": { + "name": "batch_process", + "description": "批量处理文件", + "parameters": { + "type": "object", + "properties": { + "file_paths": { + "type": "array", + "description": "文件路径列表", + "items": {"type": "string"} + }, + "priorities": { + "type": "array", + "description": "优先级列表", + "items": {"type": "integer"} + } + }, + "required": ["file_paths"] + } + } + } + + # 解析 + tool_call = ToolCall.model_validate(raw_definition) + print(f"工具名称: {tool_call.name}") + print(f"必填参数: {tool_call.input_required}") + + file_paths_attr = tool_call.input_schema["file_paths"] + print(f"file_paths 类型: {file_paths_attr.type}") + print( + f"file_paths items 类型: {file_paths_attr.items.type if hasattr(file_paths_attr.items, 'type') else file_paths_attr.items}") + + # 导出并验证相等性 + dumped_data = tool_call.simple_input_dump() + print(f"\n原始定义:\n{json.dumps(raw_definition, indent=2, ensure_ascii=False)}") + print(f"\n导出结果:\n{json.dumps(dumped_data, indent=2, ensure_ascii=False)}") + + # 验证相等 + assert dumped_data == raw_definition, "基本类型数组 Schema 导出结果与原始定义不一致" + print("\n✅ 基本类型数组 Schema 测试通过：raw_definition == simple_input_dump()") + + +def test_deep_nested_schema(): + """测试深层嵌套：三层以上的嵌套结构""" + print("\n========== 测试深层嵌套 Schema ==========") + + raw_definition = { + "type": "function", + "function": { + "name": "create_project", + "description": "创建项目，包含复杂的团队和任务结构", + "parameters": { + "type": "object", + "properties": { + "project_name": {"type": "string", "description": "项目名称"}, + "team": { + "type": "object", + "description": "团队信息", + "properties": { + "leader": { + "type": "object", + "description": "团队负责人", + "properties": { + "name": {"type": "string"}, + "contact": { + "type": "object", + "properties": { + "email": {"type": "string"}, + "phone": {"type": "string"} + }, + "required": ["email"] + } + }, + "required": ["name", "contact"] + }, + "members": { + "type": "array", + "description": "团队成员列表", + "items": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "role": {"type": "string"}, + "skills": { + "type": "array", + "items": {"type": "string"} + } + }, + "required": ["name", "role"] + } + } + }, + "required": ["leader"] + } + }, + "required": ["project_name", "team"] + } + } + } + + # 解析 + tool_call = ToolCall.model_validate(raw_definition) + print(f"工具名称: {tool_call.name}") + print(f"根级必填项: {tool_call.input_required}") + + team_attr = tool_call.input_schema["team"] + leader_attr = team_attr.properties["leader"] + contact_attr = leader_attr.properties["contact"] + print(f"Team 必填项: {team_attr.required}") + print(f"Leader 必填项: {leader_attr.required}") + print(f"Contact 必填项: {contact_attr.required}") + + # 导出并验证相等性 + dumped_data = tool_call.simple_input_dump() + print(f"\n原始定义:\n{json.dumps(raw_definition, indent=2, ensure_ascii=False)}") + print(f"\n导出结果:\n{json.dumps(dumped_data, indent=2, ensure_ascii=False)}") + + # 验证相等 + assert dumped_data == raw_definition, "深层嵌套 Schema 导出结果与原始定义不一致" + print("\n✅ 深层嵌套 Schema 测试通过：raw_definition == simple_input_dump()") + + +def test_mixed_types_schema(): + """测试混合类型：包含所有基本类型和嵌套类型""" + print("\n========== 测试混合类型 Schema ==========") + + raw_definition = { + "type": "function", + "function": { + "name": "configure_system", + "description": "配置系统参数，包含各种类型", + "parameters": { + "type": "object", + "properties": { + "enabled": {"type": "boolean", "description": "是否启用"}, + "max_connections": {"type": "integer", "description": "最大连接数"}, + "timeout": {"type": "number", "description": "超时时间（秒）"}, + "mode": { + "type": "string", + "description": "运行模式", + "enum": ["development", "production", "testing"] + }, + "allowed_ips": { + "type": "array", + "description": "允许的IP地址列表", + "items": {"type": "string"} + }, + "database": { + "type": "object", + "description": "数据库配置", + "properties": { + "host": {"type": "string"}, + "port": {"type": "integer"}, + "ssl_enabled": {"type": "boolean"} + }, + "required": ["host", "port"] + } + }, + "required": ["enabled", "mode"] + } + } + } + + # 解析 + tool_call = ToolCall.model_validate(raw_definition) + print(f"工具名称: {tool_call.name}") + print(f"根级必填项: {tool_call.input_required}") + + # 验证各种类型 + print(f"enabled 类型: {tool_call.input_schema['enabled'].type}") + print(f"max_connections 类型: {tool_call.input_schema['max_connections'].type}") + print(f"timeout 类型: {tool_call.input_schema['timeout'].type}") + print(f"mode 枚举值: {tool_call.input_schema['mode'].enum}") + + # 导出并验证相等性 + dumped_data = tool_call.simple_input_dump() + print(f"\n原始定义:\n{json.dumps(raw_definition, indent=2, ensure_ascii=False)}") + print(f"\n导出结果:\n{json.dumps(dumped_data, indent=2, ensure_ascii=False)}") + + # 验证相等 + assert dumped_data == raw_definition, "混合类型 Schema 导出结果与原始定义不一致" + print("\n✅ 混合类型 Schema 测试通过：raw_definition == simple_input_dump()") + + +if __name__ == "__main__": + test_simple_schema() + test_medium_nested_schema() + test_nested_schema() + test_array_of_primitives() + test_deep_nested_schema() + test_mixed_types_schema() + print("\n" + "=" * 50) + print("🎉 所有测试用例通过！") + print("=" * 50) From a7301f99aed8ffde5b744e608a17c21972614034 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 00:01:26 +0800 Subject: [PATCH 05/10] feat(context): add comprehensive context management system with prompt handling and registries --- reme_ai/core/context/__init__.py | 16 ++++ reme_ai/core/context/base_context.py | 41 +++++++++ reme_ai/core/context/prompt_handler.py | 95 +++++++++++++++++++++ reme_ai/core/context/registry.py | 19 +++++ reme_ai/core/context/runtime_context.py | 56 +++++++++++++ reme_ai/core/context/service_context.py | 105 ++++++++++++++++++++++++ reme_ai/core/schema/tool_call.py | 12 +-- tests/test_message.py | 6 +- tests/test_tool_call.py | 100 +++++++++++----------- 9 files changed, 393 insertions(+), 57 deletions(-) create mode 100644 reme_ai/core/context/__init__.py create mode 100644 reme_ai/core/context/base_context.py create mode 100644 reme_ai/core/context/prompt_handler.py create mode 100644 reme_ai/core/context/registry.py create mode 100644 reme_ai/core/context/runtime_context.py create mode 100644 reme_ai/core/context/service_context.py diff --git a/reme_ai/core/context/__init__.py b/reme_ai/core/context/__init__.py new file mode 100644 index 00000000..7f26d600 --- /dev/null +++ b/reme_ai/core/context/__init__.py @@ -0,0 +1,16 @@ +"""context""" + +from .base_context import BaseContext +from .prompt_handler import PromptHandler +from .registry import Registry +from .runtime_context import RuntimeContext +from .service_context import ServiceContext, C + +__all__ = [ + "BaseContext", + "PromptHandler", + "Registry", + "RuntimeContext", + "ServiceContext", + "C", +] diff --git a/reme_ai/core/context/base_context.py b/reme_ai/core/context/base_context.py new file mode 100644 index 00000000..dabd8cdb --- /dev/null +++ b/reme_ai/core/context/base_context.py @@ -0,0 +1,41 @@ +"""Module providing a dictionary subclass with attribute-style access and pickling support.""" + +from typing import Generic, TypeVar + +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") + + +class BaseContext(dict, Generic[_KT, _VT]): + """A dictionary subclass that enables accessing and modifying keys as attributes.""" + + def __getattr__(self, name: str) -> _VT: + """Retrieve a dictionary item as an attribute.""" + try: + return self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __setattr__(self, name: str, value: _VT) -> None: + """Assign a value to a dictionary item using attribute syntax.""" + self[name] = value + + def __delattr__(self, name: str) -> None: + """Remove a dictionary item using attribute syntax.""" + try: + # Delete item from dict via key + del self[name] + except KeyError as e: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e + + def __getstate__(self) -> dict: + """Return the dictionary representation for pickling.""" + return dict(self) + + def __setstate__(self, state: dict) -> None: + """Restore the dictionary state from a pickled object.""" + self.update(state) + + def __reduce__(self): + """Define the reconstruction logic for pickling processes.""" + return self.__class__, (), self.__getstate__() diff --git a/reme_ai/core/context/prompt_handler.py b/reme_ai/core/context/prompt_handler.py new file mode 100644 index 00000000..e48f98ac --- /dev/null +++ b/reme_ai/core/context/prompt_handler.py @@ -0,0 +1,95 @@ +"""Module for managing and formatting prompt templates from files or dictionaries.""" + +from pathlib import Path + +import yaml +from loguru import logger + +from .base_context import BaseContext +from .service_context import C + + +class PromptHandler(BaseContext): + """A context-aware handler for loading, retrieving, and formatting prompt templates.""" + + def __init__(self, language: str = "", **kwargs): + """Initialize the handler with a specific language and optional context data.""" + super().__init__(**kwargs) + self.language: str = language or C.language + + def load_prompt_by_file(self, prompt_file_path: Path | str = None): + """Load prompt configurations from a YAML file into the context.""" + if prompt_file_path is None: + return self + + if isinstance(prompt_file_path, str): + prompt_file_path = Path(prompt_file_path) + + if not prompt_file_path.exists(): + return self + + with prompt_file_path.open(encoding="utf-8") as f: + # Load YAML content using the full loader + prompt_dict = yaml.load(f, yaml.FullLoader) + self.load_prompt_dict(prompt_dict) + return self + + def load_prompt_dict(self, prompt_dict: dict = None): + """Merge a dictionary of prompt strings into the current context.""" + if not prompt_dict: + return self + + for key, value in prompt_dict.items(): + if isinstance(value, str): + if key in self: + logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}") + else: + logger.debug(f"Adding new prompt key={key}, value={value}") + self[key] = value + return self + + def get_prompt(self, prompt_name: str): + """Retrieve a prompt by name, automatically appending the language suffix if needed.""" + key: str = prompt_name + if self.language and not key.endswith(self.language.strip()): + key += "_" + self.language.strip() + + assert key in self, f"prompt_name={key} not found." + return self[key] + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Format a prompt by filtering flagged lines and filling template variables.""" + prompt = self.get_prompt(prompt_name) + + # Separate boolean flags from string formatting arguments + flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + + if flag_kwargs: + split_prompt = [] + for line in prompt.strip().split("\n"): + hit = False + hit_flag = True + for key, flag in flag_kwargs.items(): + if not line.startswith(f"[{key}]"): + continue + + hit = True + hit_flag = flag + # Remove the flag prefix from the line + line = line.strip(f"[{key}]") + break + + # Include line if no flag is present or if the flag evaluates to True + if not hit: + split_prompt.append(line) + elif hit_flag: + split_prompt.append(line) + + prompt = "\n".join(split_prompt) + + if other_kwargs: + # Apply standard Python string formatting + prompt = prompt.format(**other_kwargs) + + return prompt diff --git a/reme_ai/core/context/registry.py b/reme_ai/core/context/registry.py new file mode 100644 index 00000000..f403037d --- /dev/null +++ b/reme_ai/core/context/registry.py @@ -0,0 +1,19 @@ +"""Module providing a registry class for managing class-to-name mappings via decorators.""" + +from .base_context import BaseContext + + +class Registry(BaseContext): + """A registry container that uses decorators to map and store class references.""" + + def register(self, name: str = "", add_cls: bool = True): + """Return a decorator that registers a class under a specific name in the registry.""" + + def decorator(cls): + if add_cls: + # Use provided name or default to the class name as the key + key = name or cls.__name__ + self[key] = cls + return cls + + return decorator diff --git a/reme_ai/core/context/runtime_context.py b/reme_ai/core/context/runtime_context.py new file mode 100644 index 00000000..ded35ded --- /dev/null +++ b/reme_ai/core/context/runtime_context.py @@ -0,0 +1,56 @@ +"""Module providing a runtime context for managing response states and asynchronous data streaming.""" + +import asyncio + +from .base_context import BaseContext +from ..enumeration import ChunkEnum +from ..schema import Response +from ..schema import StreamChunk + + +class RuntimeContext(BaseContext): + """A context class for handling execution state, including response metadata and stream queues.""" + + def __init__( + self, + response: Response | None = None, + stream_queue: asyncio.Queue | None = None, + **kwargs, + ): + """Initialize the runtime context with optional response objects and message queues.""" + super().__init__(**kwargs) + + self.response: Response | None = response if response is not None else Response() + self.stream_queue: asyncio.Queue | None = stream_queue + + async def add_stream_string_and_type(self, chunk: str, chunk_type: ChunkEnum): + """Create and enqueue a stream chunk from a raw string and specific type.""" + if self.stream_queue is None: + return self + + # Package raw data into a StreamChunk schema + stream_chunk = StreamChunk(chunk_type=chunk_type, chunk=chunk) + await self.stream_queue.put(stream_chunk) + return self + + async def add_stream_chunk(self, stream_chunk: StreamChunk): + """Directly enqueue an existing stream chunk into the stream queue.""" + if self.stream_queue is None: + return self + await self.stream_queue.put(stream_chunk) + return self + + async def add_stream_done(self): + """Enqueue a termination chunk to signal the end of the data stream.""" + if self.stream_queue is None: + return self + + # Create a special chunk representing the completion state + done_chunk = StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True) + await self.stream_queue.put(done_chunk) + return self + + def add_response_error(self, e: Exception): + """Update the internal response object to reflect a failure state using exception details.""" + self.response.success = False + self.response.answer = str(e.args) diff --git a/reme_ai/core/context/service_context.py b/reme_ai/core/context/service_context.py new file mode 100644 index 00000000..af99a870 --- /dev/null +++ b/reme_ai/core/context/service_context.py @@ -0,0 +1,105 @@ +"""Module for managing global service configurations and component registries via a singleton context.""" + +from concurrent.futures import ThreadPoolExecutor +from typing import Dict + +from .base_context import BaseContext +from .registry import Registry +from ..enumeration import RegistryEnum +from ..schema import ServiceConfig +from ..utils import singleton + + +@singleton +class ServiceContext(BaseContext): + """A singleton container for global application state, thread pools, and component registries.""" + + def __init__(self, **kwargs): + """Initialize the global context with configuration objects and specialized registries.""" + super().__init__(**kwargs) + + self.service_config: ServiceConfig | None = None + self.language: str = "" + self.thread_pool: ThreadPoolExecutor | None = None + self.vector_store_dict: Dict[str, dict] = {} + self.external_mcp_tool_call_dict: dict = {} + # Initialize a registry for every category defined in RegistryEnum + self.registry_dict: Dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()} + self.flow_dict: dict = {} + + def register(self, name: str, register_type: RegistryEnum): + """Return a decorator to register a component within a specific registry category.""" + return self.registry_dict[register_type].register(name=name) + + def register_llm(self, name: str = ""): + """Register a Large Language Model class.""" + return self.register(name=name, register_type=RegistryEnum.LLM) + + def register_embedding_model(self, name: str = ""): + """Register an embedding model class.""" + return self.register(name=name, register_type=RegistryEnum.EMBEDDING_MODEL) + + def register_vector_store(self, name: str = ""): + """Register a vector store implementation class.""" + return self.register(name=name, register_type=RegistryEnum.VECTOR_STORE) + + def register_op(self, name: str = ""): + """Register an operation (Op) class.""" + return self.register(name=name, register_type=RegistryEnum.OP) + + def register_flow(self, name: str = ""): + """Register a workflow or logic flow class.""" + return self.register(name=name, register_type=RegistryEnum.FLOW) + + def register_service(self, name: str = ""): + """Register a backend service class.""" + return self.register(name=name, register_type=RegistryEnum.SERVICE) + + def register_token_counter(self, name: str = ""): + """Register a token counting utility class.""" + return self.register(name=name, register_type=RegistryEnum.TOKEN_COUNTER) + + def get_model_class(self, name: str, register_type: RegistryEnum): + """Retrieve a registered class by name from a specific registry category.""" + assert name in self.registry_dict[register_type], f"{name} not in registry_dict[{register_type}]" + return self.registry_dict[register_type][name] + + def get_embedding_model_class(self, name: str): + """Get the embedding model class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL) + + def get_llm_class(self, name: str): + """Get the LLM class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.LLM) + + def get_vector_store_class(self, name: str): + """Get the vector store class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.VECTOR_STORE) + + def get_op_class(self, name: str): + """Get the operation class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.OP) + + def get_flow_class(self, name: str): + """Get the flow class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.FLOW) + + def get_service_class(self, name: str): + """Get the service class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.SERVICE) + + def get_token_counter_class(self, name: str): + """Get the token counter class registered under the given name.""" + return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER) + + def get_vector_store(self, name: str): + """Retrieve a specific vector store instance by name.""" + return self.vector_store_dict[name] + + def get_flow(self, name: str): + """Retrieve a specific flow instance by name.""" + return self.flow_dict[name] + + +# Export a global instance for easy access across the application +C = ServiceContext() diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index 3688da54..d190974e 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core/schema/tool_call.py @@ -1,6 +1,7 @@ """ MCP Tool Schema definitions for recursive JSON Schema representation. """ + import json from typing import Any, Dict, List, Literal, Optional, Union @@ -12,6 +13,7 @@ from ..enumeration.json_schema_enum import JsonSchemaEnum class ToolAttr(BaseModel): """Recursive model representing JSON Schema attributes for tool parameters.""" + model_config = ConfigDict(extra="allow") type: Literal[ @@ -23,8 +25,8 @@ class ToolAttr(BaseModel): JsonSchemaEnum.BOOLEAN.value, JsonSchemaEnum.NULL.value, ] = Field( - default=JsonSchemaEnum.STRING.value, - description="The data type of the attribute" + default=JsonSchemaEnum.STRING.value, + description="The data type of the attribute", ) description: Optional[str] = Field(default=None, description="Description of the attribute") required: Optional[List[str]] = Field(default=None, description="Required property names for object types") @@ -32,7 +34,6 @@ class ToolAttr(BaseModel): items: Optional[Union[Dict[str, Any], "ToolAttr"]] = Field(default=None, description="Schema for array items") enum: Optional[List[str]] = Field(default=None, description="Allowed values for the attribute") - def simple_input_dump(self) -> dict: """Serializes the attribute into a standard JSON Schema dictionary.""" res: dict = {"type": self.type} @@ -42,8 +43,9 @@ class ToolAttr(BaseModel): res["enum"] = self.enum if self.type == "object" and self.properties: - res["properties"] = {k: v.simple_input_dump() if isinstance(v, ToolAttr) else v - for k, v in self.properties.items()} + res["properties"] = { + k: v.simple_input_dump() if isinstance(v, ToolAttr) else v for k, v in self.properties.items() + } if self.required: res["required"] = self.required diff --git a/tests/test_message.py b/tests/test_message.py index 35d8598e..b45563dd 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -27,7 +27,7 @@ class TestModelDefinitions(unittest.TestCase): self.assertEqual(dump["type"], "string") self.assertEqual(dump["enum"], ["Beijing", "London"]) self.assertIn("description", dump) - + # Test object attribute with required child properties obj_attr = ToolAttr( type="object", @@ -39,10 +39,10 @@ class TestModelDefinitions(unittest.TestCase): required=["name"], # 'name' is required, 'age' is optional ) obj_dump = obj_attr.simple_input_dump() - + print("\n=== ToolAttr.simple_input_dump() (object with required) ===") print(obj_dump) - + self.assertEqual(obj_dump["type"], "object") self.assertIn("properties", obj_dump) self.assertEqual(obj_dump["required"], ["name"]) diff --git a/tests/test_tool_call.py b/tests/test_tool_call.py index 242f07ee..f914dca7 100644 --- a/tests/test_tool_call.py +++ b/tests/test_tool_call.py @@ -1,3 +1,5 @@ +"""simple tool call test""" + import json from reme_ai.core.schema.tool_call import ToolCall @@ -16,11 +18,11 @@ def test_simple_schema(): "type": "object", "properties": { "city": {"type": "string", "description": "城市名称"}, - "unit": {"type": "string", "description": "温度单位", "enum": ["celsius", "fahrenheit"]} + "unit": {"type": "string", "description": "温度单位", "enum": ["celsius", "fahrenheit"]}, }, - "required": ["city"] - } - } + "required": ["city"], + }, + }, } # 解析 @@ -58,14 +60,14 @@ def test_medium_nested_schema(): "properties": { "name": {"type": "string", "description": "客户姓名"}, "email": {"type": "string", "description": "客户邮箱"}, - "phone": {"type": "string", "description": "联系电话"} + "phone": {"type": "string", "description": "联系电话"}, }, - "required": ["name", "email"] - } + "required": ["name", "email"], + }, }, - "required": ["order_id", "customer"] - } - } + "required": ["order_id", "customer"], + }, + }, } # 解析 @@ -106,9 +108,9 @@ def test_nested_schema(): "description": "用户元数据", "properties": { "age": {"type": "integer"}, - "location": {"type": "string"} + "location": {"type": "string"}, }, - "required": ["age"] + "required": ["age"], }, "tags": { "type": "array", @@ -117,15 +119,15 @@ def test_nested_schema(): "type": "object", "properties": { "tag_id": {"type": "string"}, - "level": {"type": "number"} + "level": {"type": "number"}, }, - "required": ["tag_id"] - } - } + "required": ["tag_id"], + }, + }, }, - "required": ["username", "metadata"] - } - } + "required": ["username", "metadata"], + }, + }, } # 2. 解析：将原始字典转化为 ToolCall 实例 @@ -171,17 +173,17 @@ def test_array_of_primitives(): "file_paths": { "type": "array", "description": "文件路径列表", - "items": {"type": "string"} + "items": {"type": "string"}, }, "priorities": { "type": "array", "description": "优先级列表", - "items": {"type": "integer"} - } + "items": {"type": "integer"}, + }, }, - "required": ["file_paths"] - } - } + "required": ["file_paths"], + }, + }, } # 解析 @@ -191,8 +193,8 @@ def test_array_of_primitives(): file_paths_attr = tool_call.input_schema["file_paths"] print(f"file_paths 类型: {file_paths_attr.type}") - print( - f"file_paths items 类型: {file_paths_attr.items.type if hasattr(file_paths_attr.items, 'type') else file_paths_attr.items}") + t_items_type = file_paths_attr.items.type if hasattr(file_paths_attr.items, "type") else file_paths_attr.items + print(f"file_paths items 类型: {t_items_type}") # 导出并验证相等性 dumped_data = tool_call.simple_input_dump() @@ -230,12 +232,12 @@ def test_deep_nested_schema(): "type": "object", "properties": { "email": {"type": "string"}, - "phone": {"type": "string"} + "phone": {"type": "string"}, }, - "required": ["email"] - } + "required": ["email"], + }, }, - "required": ["name", "contact"] + "required": ["name", "contact"], }, "members": { "type": "array", @@ -247,19 +249,19 @@ def test_deep_nested_schema(): "role": {"type": "string"}, "skills": { "type": "array", - "items": {"type": "string"} - } + "items": {"type": "string"}, + }, }, - "required": ["name", "role"] - } - } + "required": ["name", "role"], + }, + }, }, - "required": ["leader"] - } + "required": ["leader"], + }, }, - "required": ["project_name", "team"] - } - } + "required": ["project_name", "team"], + }, + }, } # 解析 @@ -302,12 +304,12 @@ def test_mixed_types_schema(): "mode": { "type": "string", "description": "运行模式", - "enum": ["development", "production", "testing"] + "enum": ["development", "production", "testing"], }, "allowed_ips": { "type": "array", "description": "允许的IP地址列表", - "items": {"type": "string"} + "items": {"type": "string"}, }, "database": { "type": "object", @@ -315,14 +317,14 @@ def test_mixed_types_schema(): "properties": { "host": {"type": "string"}, "port": {"type": "integer"}, - "ssl_enabled": {"type": "boolean"} + "ssl_enabled": {"type": "boolean"}, }, - "required": ["host", "port"] - } + "required": ["host", "port"], + }, }, - "required": ["enabled", "mode"] - } - } + "required": ["enabled", "mode"], + }, + }, } # 解析 From 457284a046c3e7f9c1493c4b8661c14980c8fdf6 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 00:49:45 +0800 Subject: [PATCH 06/10] feat(llm): add LLM module with base interface and multiple provider implementations --- reme_ai/core/llm/__init__.py | 15 + reme_ai/core/llm/base_llm.py | 364 ++++++++++++++++++++++++ reme_ai/core/llm/lite_llm.py | 118 ++++++++ reme_ai/core/llm/lite_llm_sync.py | 101 +++++++ reme_ai/core/llm/openai_llm.py | 117 ++++++++ reme_ai/core/llm/openai_llm_sync.py | 71 +++++ reme_ai/core/schema/tool_call.py | 76 ++--- reme_ai/core/utils/__init__.py | 12 +- reme_ai/core/utils/env_utils.py | 63 +++++ tests/test_llm.py | 422 ++++++++++++++++++++++++++++ tests/test_llm_sync.py | 421 +++++++++++++++++++++++++++ tests/test_message.py | 16 +- tests/test_tool_call.py | 28 +- 13 files changed, 1766 insertions(+), 58 deletions(-) create mode 100644 reme_ai/core/llm/__init__.py create mode 100644 reme_ai/core/llm/base_llm.py create mode 100644 reme_ai/core/llm/lite_llm.py create mode 100644 reme_ai/core/llm/lite_llm_sync.py create mode 100644 reme_ai/core/llm/openai_llm.py create mode 100644 reme_ai/core/llm/openai_llm_sync.py create mode 100644 reme_ai/core/utils/env_utils.py create mode 100644 tests/test_llm.py create mode 100644 tests/test_llm_sync.py diff --git a/reme_ai/core/llm/__init__.py b/reme_ai/core/llm/__init__.py new file mode 100644 index 00000000..57578f49 --- /dev/null +++ b/reme_ai/core/llm/__init__.py @@ -0,0 +1,15 @@ +"""llm""" + +from .base_llm import BaseLLM +from .lite_llm import LiteLLM +from .lite_llm_sync import LiteLLMSync +from .openai_llm import OpenAILLM +from .openai_llm_sync import OpenAILLMSync + +__all__ = [ + "BaseLLM", + "LiteLLM", + "LiteLLMSync", + "OpenAILLM", + "OpenAILLMSync", +] diff --git a/reme_ai/core/llm/base_llm.py b/reme_ai/core/llm/base_llm.py new file mode 100644 index 00000000..6f1da121 --- /dev/null +++ b/reme_ai/core/llm/base_llm.py @@ -0,0 +1,364 @@ +"""Abstract base interface for ReMe LLM implementations.""" + +import asyncio +import json +import time +from abc import ABC +from typing import List, Callable, Generator, AsyncGenerator, Any, Optional, Dict + +from loguru import logger + +from ..enumeration import ChunkEnum, Role +from ..schema import Message +from ..schema import StreamChunk +from ..schema import ToolCall + + +class BaseLLM(ABC): + """Abstract base class defining the standard interface for LLM interactions.""" + + def __init__(self, model_name: str, max_retries: int = 3, raise_exception: bool = False, **kwargs): + """Initialize the LLM client with model configurations and retry policies.""" + self.model_name: str = model_name + self.max_retries: int = max_retries + self.raise_exception: bool = raise_exception + self.kwargs: dict = kwargs + + @staticmethod + def _process_stream_chunk( + stream_chunk: StreamChunk, + state: dict, + enable_stream_print: bool = False, + ) -> None: + """Update the aggregation state by processing an individual stream chunk.""" + if stream_chunk.chunk_type is ChunkEnum.USAGE: + if enable_stream_print: + print(f"\n{json.dumps(stream_chunk.chunk, ensure_ascii=False, indent=2)}", flush=True) + + elif stream_chunk.chunk_type is ChunkEnum.THINK: + if enable_stream_print: + if not state["enter_think"]: + state["enter_think"] = True + print("\n", end="", flush=True) + print(stream_chunk.chunk, end="", flush=True) + state["reasoning_content"] += stream_chunk.chunk + + elif stream_chunk.chunk_type is ChunkEnum.ANSWER: + if enable_stream_print: + if not state["enter_answer"]: + state["enter_answer"] = True + if state["enter_think"]: + print("\n", flush=True) + print(stream_chunk.chunk, end="", flush=True) + state["answer_content"] += stream_chunk.chunk + + elif stream_chunk.chunk_type is ChunkEnum.TOOL: + if enable_stream_print: + print(f"\n{json.dumps(stream_chunk.chunk, ensure_ascii=False, indent=2)}", flush=True) + state["tool_calls"].append(stream_chunk.chunk) + + elif stream_chunk.chunk_type is ChunkEnum.ERROR: + if enable_stream_print: + print(f"\n{stream_chunk.chunk}", flush=True) + + @staticmethod + def _create_message_from_state(state: dict) -> Message: + """Construct a Message object from the accumulated aggregation state.""" + return Message( + role=Role.ASSISTANT, + reasoning_content=state["reasoning_content"], + content=state["answer_content"], + tool_calls=state["tool_calls"], + ) + + @staticmethod + def _accumulate_tool_call_chunk( + tool_call, + ret_tools: List[ToolCall], + ) -> None: + """Assemble incremental tool call fragments into complete ToolCall objects.""" + index = tool_call.index + + # Ensure we have a ToolCall object at this index + while len(ret_tools) <= index: + ret_tools.append(ToolCall(index=index)) + + # Accumulate tool call parts (id, name, arguments) + if tool_call.id: + ret_tools[index].id += tool_call.id + + if tool_call.function and tool_call.function.name: + ret_tools[index].name += tool_call.function.name + + if tool_call.function and tool_call.function.arguments: + ret_tools[index].arguments += tool_call.function.arguments + + @staticmethod + def _validate_and_serialize_tools( + ret_tools: List[ToolCall], + tools: Optional[List[ToolCall]], + ) -> List[Dict]: + """Validate tool call integrity and return serialized tool dictionaries.""" + if not ret_tools: + return [] + + # Create lookup dict for tool validation + tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {} + validated_tools = [] + + for tool in ret_tools: + # Skip tools that weren't in the provided tool list + if tool.name not in tool_dict: + continue + + # Validate tool arguments are valid JSON + if not tool.check_argument(): + raise ValueError( + f"Tool call {tool.name} has invalid JSON arguments: {tool.arguments}", + ) + + validated_tools.append(tool.simple_output_dump()) + + return validated_tools + + def _build_stream_kwargs( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + log_params: bool = True, + **kwargs, + ) -> dict: + """Construct provider-specific parameters for streaming API requests.""" + raise NotImplementedError + + async def _stream_chat( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + stream_kwargs: Optional[dict] = None, + ) -> AsyncGenerator[StreamChunk, None]: + """Internal async generator for streaming raw response chunks.""" + raise NotImplementedError + + def _stream_chat_sync( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + stream_kwargs: Optional[dict] = None, + ) -> Generator[StreamChunk, None, None]: + """Internal synchronous generator for streaming raw response chunks.""" + raise NotImplementedError + + async def _stream_with_retry( + self, + operation_name: str, + messages: List[Message], + tools: Optional[List[ToolCall]], + stream_kwargs: dict, + ) -> AsyncGenerator[StreamChunk, None]: + """Execute the async streaming operation with retry logic and error recovery.""" + for i in range(self.max_retries): + try: + async for chunk in self._stream_chat(messages=messages, tools=tools, stream_kwargs=stream_kwargs): + yield chunk + return + + except Exception as e: + logger.exception(f"{operation_name} with model={self.model_name} encounter error with e={e.args}") + + if i == self.max_retries - 1: + if self.raise_exception: + raise e + yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) + return + + yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) + await asyncio.sleep(i + 1) + + def _stream_with_retry_sync( + self, + operation_name: str, + messages: List[Message], + tools: Optional[List[ToolCall]], + stream_kwargs: dict, + ) -> Generator[StreamChunk, None, None]: + """Execute the synchronous streaming operation with retry logic and error recovery.""" + for i in range(self.max_retries): + try: + yield from self._stream_chat_sync(messages=messages, tools=tools, stream_kwargs=stream_kwargs) + return + + except Exception as e: + logger.exception(f"{operation_name} with model={self.model_name} encounter error with e={e.args}") + + if i == self.max_retries - 1: + if self.raise_exception: + raise e + yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) + return + + yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) + time.sleep(i + 1) + + async def stream_chat( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + **kwargs, + ) -> AsyncGenerator[StreamChunk, None]: + """Public async interface for streaming chat completions with retries.""" + stream_kwargs = self._build_stream_kwargs(messages, tools, **kwargs) + async for chunk in self._stream_with_retry("stream chat", messages, tools, stream_kwargs): + yield chunk + + def stream_chat_sync( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + **kwargs, + ) -> Generator[StreamChunk, None, None]: + """Public synchronous interface for streaming chat completions with retries.""" + stream_kwargs = self._build_stream_kwargs(messages, tools, **kwargs) + yield from self._stream_with_retry_sync("stream chat sync", messages, tools, stream_kwargs) + + async def _chat( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + enable_stream_print: bool = False, + **kwargs, + ) -> Message: + """Internal async method to aggregate a full response by consuming the stream.""" + state = { + "enter_think": False, + "enter_answer": False, + "reasoning_content": "", + "answer_content": "", + "tool_calls": [], + } + + stream_kwargs = self._build_stream_kwargs(messages, tools, **kwargs) + async for stream_chunk in self._stream_chat(messages=messages, tools=tools, stream_kwargs=stream_kwargs): + self._process_stream_chunk(stream_chunk, state, enable_stream_print) + + return self._create_message_from_state(state) + + def _chat_sync( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + enable_stream_print: bool = False, + **kwargs, + ) -> Message: + """Internal synchronous method to aggregate a full response by consuming the stream.""" + state = { + "enter_think": False, + "enter_answer": False, + "reasoning_content": "", + "answer_content": "", + "tool_calls": [], + } + + stream_kwargs = self._build_stream_kwargs(messages, tools, **kwargs) + for stream_chunk in self._stream_chat_sync(messages=messages, tools=tools, stream_kwargs=stream_kwargs): + self._process_stream_chunk(stream_chunk, state, enable_stream_print) + + return self._create_message_from_state(state) + + async def _execute_with_retry( + self, + operation_name: str, + operation_fn: Callable[[], Any], + callback_fn: Optional[Callable[[Message], Any]] = None, + default_value: Any = None, + ) -> Message | Any: + """Execute a generic async operation with error handling and retry logic.""" + for i in range(self.max_retries): + try: + result = await operation_fn() + return callback_fn(result) if callback_fn else result + + except Exception as e: + logger.exception(f"{operation_name} with model={self.model_name} encounter error with e={e.args}") + + if i == self.max_retries - 1: + if self.raise_exception: + raise e + return default_value + + await asyncio.sleep(1 + i) + return default_value + + def _execute_with_retry_sync( + self, + operation_name: str, + operation_fn: Callable[[], Message], + callback_fn: Optional[Callable[[Message], Any]] = None, + default_value: Any = None, + ) -> Message | Any: + """Execute a generic synchronous operation with error handling and retry logic.""" + for i in range(self.max_retries): + try: + result = operation_fn() + return callback_fn(result) if callback_fn else result + + except Exception as e: + logger.exception(f"{operation_name} with model={self.model_name} encounter error with e={e.args}") + + if i == self.max_retries - 1: + if self.raise_exception: + raise e + return default_value + + time.sleep(1 + i) + return default_value + + async def chat( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + enable_stream_print: bool = False, + callback_fn: Optional[Callable[[Message], Any]] = None, + default_value: Any = None, + **kwargs, + ) -> Message | Any: + """Perform an async chat completion with integrated retries and error handling.""" + return await self._execute_with_retry( + operation_name="chat", + operation_fn=lambda: self._chat( + messages=messages, + tools=tools, + enable_stream_print=enable_stream_print, + **kwargs, + ), + callback_fn=callback_fn, + default_value=default_value, + ) + + def chat_sync( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + enable_stream_print: bool = False, + callback_fn: Optional[Callable[[Message], Any]] = None, + default_value: Any = None, + **kwargs, + ) -> Message | Any: + """Perform a synchronous chat completion with integrated retries and error handling.""" + return self._execute_with_retry_sync( + operation_name="chat sync", + operation_fn=lambda: self._chat_sync( + messages=messages, + tools=tools, + enable_stream_print=enable_stream_print, + **kwargs, + ), + callback_fn=callback_fn, + default_value=default_value, + ) + + async def close(self): + """Release any asynchronous resources or connections held by the client.""" + + def close_sync(self): + """Release any synchronous resources or connections held by the client.""" diff --git a/reme_ai/core/llm/lite_llm.py b/reme_ai/core/llm/lite_llm.py new file mode 100644 index 00000000..32d3005c --- /dev/null +++ b/reme_ai/core/llm/lite_llm.py @@ -0,0 +1,118 @@ +"""LiteLLM asynchronous implementation for ReMe.""" + +import os +from typing import List, AsyncGenerator, Optional + +import litellm +from loguru import logger + +from .base_llm import BaseLLM +from ..context import C +from ..enumeration import ChunkEnum +from ..schema import Message +from ..schema import StreamChunk +from ..schema import ToolCall + + +@C.register_llm("litellm") +class LiteLLM(BaseLLM): + """Async LLM implementation using LiteLLM to support multiple providers.""" + + def __init__( + self, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + custom_llm_provider: str = "openai", + **kwargs, + ): + """Initialize the LiteLLM client with API configuration and provider settings.""" + super().__init__(**kwargs) + self.api_key: Optional[str] = api_key or os.getenv("REME_LLM_API_KEY") + self.base_url: Optional[str] = base_url or os.getenv("REME_LLM_BASE_URL") + self.custom_llm_provider: str = custom_llm_provider + + def _build_stream_kwargs( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + log_params: bool = True, + **kwargs, + ) -> dict: + """Construct and log the parameters dictionary for LiteLLM API calls.""" + # Construct the API parameters by merging multiple sources + llm_kwargs = { + "model": self.model_name, + "messages": [x.simple_dump() for x in messages], + "tools": [x.simple_input_dump() for x in tools] if tools else None, + "stream": True, + "custom_llm_provider": self.custom_llm_provider, + **self.kwargs, + **kwargs, + } + + # Add API key and base URL if provided + if self.api_key: + llm_kwargs["api_key"] = self.api_key + if self.base_url: + llm_kwargs["base_url"] = self.base_url + + # Log parameters for debugging, with message/tool counts instead of full content + if log_params: + log_kwargs: dict = {} + for k, v in llm_kwargs.items(): + if k in ["messages", "tools"]: + log_kwargs[k] = len(v) if v is not None else 0 + elif k == "api_key": + # Mask API key in logs for security + log_kwargs[k] = "***" if v else None + else: + log_kwargs[k] = v + logger.info(f"llm_kwargs={log_kwargs}") + + return llm_kwargs + + async def _stream_chat( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + stream_kwargs: Optional[dict] = None, + ) -> AsyncGenerator[StreamChunk, None]: + """Execute async streaming chat requests and yield processed response chunks.""" + # Create streaming completion request using LiteLLM asynchronously + stream_kwargs = stream_kwargs or {} + completion = await litellm.acompletion(**stream_kwargs) + + # Track accumulated tool calls across chunks + ret_tools: List[ToolCall] = [] + # Flag to track if we've started receiving answer content + is_answering: bool = False + + async for chunk in completion: + # Handle usage information (typically the last chunk) + if not chunk.choices: + if hasattr(chunk, "usage") and chunk.usage: + yield StreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) + + else: + delta = chunk.choices[0].delta + + # Check for reasoning content (models that support thinking) + if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: + yield StreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) + + else: + if not is_answering: + is_answering = True + + # Yield regular text content + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + + # Process tool calls - LiteLLM streams them incrementally + if hasattr(delta, "tool_calls") and delta.tool_calls is not None: + for tool_call in delta.tool_calls: + self._accumulate_tool_call_chunk(tool_call, ret_tools) + + # After streaming completes, validate and yield complete tool calls + for tool_data in self._validate_and_serialize_tools(ret_tools, tools): + yield StreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool_data) diff --git a/reme_ai/core/llm/lite_llm_sync.py b/reme_ai/core/llm/lite_llm_sync.py new file mode 100644 index 00000000..5f814a2e --- /dev/null +++ b/reme_ai/core/llm/lite_llm_sync.py @@ -0,0 +1,101 @@ +"""Synchronous LiteLLM-based LLM implementation for the ReMe framework. + +This module provides a unified synchronous interface for 100+ LLM providers via LiteLLM, +supporting streaming completions, tool calling, and reasoning content. For +asynchronous operations, refer to the LiteLLM class in the lite_llm module. +""" + +from typing import List, Generator, Optional + +import litellm + +from .lite_llm import LiteLLM +from ..context import C +from ..enumeration import ChunkEnum +from ..schema import Message +from ..schema import StreamChunk +from ..schema import ToolCall + + +@C.register_llm("litellm_sync") +class LiteLLMSync(LiteLLM): + """ + Synchronous LiteLLM client for executing chat completions and streaming responses. + + This class extends the base LiteLLM implementation to provide synchronous + execution of streaming methods, inheriting initialization and configuration + logic from the parent class. + + Example: + >>> llm = LiteLLMSync( + ... model_name="qwen3-max", + ... api_key="sk-...", + ... temperature=0.7 + ... ) + >>> messages = [Message(role=Role.USER, content="Hello!")] + >>> for chunk in llm.chat(messages): + ... print(chunk) + """ + + def _stream_chat_sync( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + stream_kwargs: Optional[dict] = None, + ) -> Generator[StreamChunk, None, None]: + """ + Internal synchronous generator for processing streaming chat completion chunks. + + This method orchestrates the LiteLLM completion lifecycle by categorizing + raw API chunks into usage data, reasoning content (thinking), regular + text responses, and aggregated tool calls. + + Args: + messages: List of conversation messages to send to the model. + tools: Optional list of tool definitions available for the model to call. + stream_kwargs: Dictionary of pre-built parameters for the LiteLLM API. + + Yields: + StreamChunk: Wrapped response fragments categorized by ChunkEnum. + + Raises: + ValueError: If tool call arguments fail validation or serialization. + """ + # Create streaming completion request using LiteLLM + stream_kwargs = stream_kwargs or {} + completion = litellm.completion(**stream_kwargs) + + # Track accumulated tool calls across chunks + ret_tools: List[ToolCall] = [] + # Flag to track if we've started receiving answer content + is_answering: bool = False + + for chunk in completion: + # Handle usage information (typically the last chunk) + if not chunk.choices: + if hasattr(chunk, "usage") and chunk.usage: + yield StreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) + + else: + delta = chunk.choices[0].delta + + # Check for reasoning content (models that support thinking) + if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: + yield StreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) + + else: + if not is_answering: + is_answering = True + + # Yield regular text content + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + + # Process tool calls - LiteLLM streams them incrementally + if hasattr(delta, "tool_calls") and delta.tool_calls is not None: + for tool_call in delta.tool_calls: + self._accumulate_tool_call_chunk(tool_call, ret_tools) + + # After streaming completes, validate and yield complete tool calls + for tool_data in self._validate_and_serialize_tools(ret_tools, tools): + yield StreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool_data) diff --git a/reme_ai/core/llm/openai_llm.py b/reme_ai/core/llm/openai_llm.py new file mode 100644 index 00000000..2c593d6f --- /dev/null +++ b/reme_ai/core/llm/openai_llm.py @@ -0,0 +1,117 @@ +"""Asynchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content.""" + +import os +from typing import List, AsyncGenerator, Optional + +from loguru import logger +from openai import AsyncOpenAI + +from .base_llm import BaseLLM +from ..context import C +from ..enumeration import ChunkEnum +from ..schema import Message +from ..schema import StreamChunk +from ..schema import ToolCall + + +@C.register_llm("openai") +class OpenAILLM(BaseLLM): + """Asynchronous LLM client for OpenAI-compatible APIs supporting streaming completions and tool execution.""" + + def __init__( + self, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + **kwargs, + ): + """Initialize the OpenAI async client with API credentials and model configuration.""" + super().__init__(**kwargs) + self.api_key: str = api_key or os.getenv("REME_LLM_API_KEY", "") + self.base_url: str = base_url or os.getenv("REME_LLM_BASE_URL", "") + + # Create client using factory method + self._client = self._create_client() + + def _create_client(self): + """Create and return an instance of the AsyncOpenAI client.""" + return AsyncOpenAI(api_key=self.api_key, base_url=self.base_url) + + def _build_stream_kwargs( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + log_params: bool = True, + **kwargs, + ) -> dict: + """Construct the parameter dictionary for the OpenAI Chat Completions API call.""" + # Construct the API parameters by merging multiple sources + llm_kwargs = { + "model": self.model_name, + "messages": [x.simple_dump() for x in messages], + "tools": [x.simple_input_dump() for x in tools] if tools else None, + "stream": True, + **self.kwargs, + **kwargs, + } + + # Log parameters for debugging, with message/tool counts instead of full content + if log_params: + log_kwargs: dict = {} + for k, v in llm_kwargs.items(): + if k in ["messages", "tools"]: + log_kwargs[k] = len(v) if v is not None else 0 + else: + log_kwargs[k] = v + logger.info(f"llm_kwargs={log_kwargs}") + + return llm_kwargs + + async def _stream_chat( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + stream_kwargs: Optional[dict] = None, + ) -> AsyncGenerator[StreamChunk, None]: + """Generate a stream of chat completion chunks including text, reasoning content, and tool calls.""" + # Create streaming completion request to OpenAI API asynchronously + stream_kwargs = stream_kwargs or {} + completion = await self._client.chat.completions.create(**stream_kwargs) + + # Track accumulated tool calls across chunks + ret_tools: List[ToolCall] = [] + # Flag to track if we've started receiving answer content + is_answering: bool = False + + async for chunk in completion: + # Handle usage information (typically the last chunk) + if not chunk.choices: + if hasattr(chunk, "usage") and chunk.usage: + yield StreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) + + else: + delta = chunk.choices[0].delta + + # Check for reasoning content (o1-preview, o1-mini models) + if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: + yield StreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) + + else: + if not is_answering: + is_answering = True + + # Yield regular text content + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + + # Process tool calls - OpenAI streams them incrementally + if delta.tool_calls is not None: + for tool_call in delta.tool_calls: + self._accumulate_tool_call_chunk(tool_call, ret_tools) + + # After streaming completes, validate and yield complete tool calls + for tool_data in self._validate_and_serialize_tools(ret_tools, tools): + yield StreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool_data) + + async def close(self): + """Asynchronously close the OpenAI client and release network resources.""" + await self._client.close() diff --git a/reme_ai/core/llm/openai_llm_sync.py b/reme_ai/core/llm/openai_llm_sync.py new file mode 100644 index 00000000..3aadab1a --- /dev/null +++ b/reme_ai/core/llm/openai_llm_sync.py @@ -0,0 +1,71 @@ +"""Synchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content.""" + +from typing import List, Generator, Optional + +from openai import OpenAI + +from .openai_llm import OpenAILLM +from ..context import C +from ..enumeration import ChunkEnum +from ..schema import Message +from ..schema import StreamChunk +from ..schema import ToolCall + + +@C.register_llm("openai_sync") +class OpenAILLMSync(OpenAILLM): + """Synchronous LLM client for OpenAI-compatible APIs, inheriting from OpenAILLM.""" + + def _create_client(self): + """Create and return an instance of the synchronous OpenAI client.""" + return OpenAI(api_key=self.api_key, base_url=self.base_url) + + def _stream_chat_sync( + self, + messages: List[Message], + tools: Optional[List[ToolCall]] = None, + stream_kwargs: Optional[dict] = None, + ) -> Generator[StreamChunk, None, None]: + """Synchronously generate a stream of chat completion chunks including text, reasoning, and tool calls.""" + # Create streaming completion request to OpenAI API + stream_kwargs = stream_kwargs or {} + completion = self._client.chat.completions.create(**stream_kwargs) + + # Track accumulated tool calls across chunks + ret_tools: List[ToolCall] = [] + # Flag to track if we've started receiving answer content + is_answering: bool = False + + for chunk in completion: + # Handle usage information (typically the last chunk) + if not chunk.choices: + if hasattr(chunk, "usage") and chunk.usage: + yield StreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) + + else: + delta = chunk.choices[0].delta + + # Check for reasoning content (o1-preview, o1-mini models) + if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: + yield StreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) + + else: + if not is_answering: + is_answering = True + + # Yield regular text content + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + + # Process tool calls - OpenAI streams them incrementally + if delta.tool_calls is not None: + for tool_call in delta.tool_calls: + self._accumulate_tool_call_chunk(tool_call, ret_tools) + + # After streaming completes, validate and yield complete tool calls + for tool_data in self._validate_and_serialize_tools(ret_tools, tools): + yield StreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool_data) + + def close_sync(self): + """Close the synchronous OpenAI client and release network resources.""" + self._client.close() diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index d190974e..e6b1092f 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core/schema/tool_call.py @@ -3,10 +3,10 @@ MCP Tool Schema definitions for recursive JSON Schema representation. """ import json -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Optional, Union from mcp.types import Tool -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator, field_validator from ..enumeration.json_schema_enum import JsonSchemaEnum @@ -16,24 +16,23 @@ class ToolAttr(BaseModel): model_config = ConfigDict(extra="allow") - type: Literal[ - JsonSchemaEnum.STRING.value, - JsonSchemaEnum.NUMBER.value, - JsonSchemaEnum.INTEGER.value, - JsonSchemaEnum.OBJECT.value, - JsonSchemaEnum.ARRAY.value, - JsonSchemaEnum.BOOLEAN.value, - JsonSchemaEnum.NULL.value, - ] = Field( - default=JsonSchemaEnum.STRING.value, - description="The data type of the attribute", - ) + type: str = Field(default=JsonSchemaEnum.STRING.value, description="The data type of the attribute") description: Optional[str] = Field(default=None, description="Description of the attribute") required: Optional[List[str]] = Field(default=None, description="Required property names for object types") properties: Optional[Dict[str, "ToolAttr"]] = Field(default=None, description="Child properties for objects") items: Optional[Union[Dict[str, Any], "ToolAttr"]] = Field(default=None, description="Schema for array items") enum: Optional[List[str]] = Field(default=None, description="Allowed values for the attribute") + @field_validator("type") + @classmethod + def validate_type_is_valid_enum(cls, v: str) -> str: + """Validates that the provided type string exists within JsonSchemaEnum values.""" + valid_types = [e.value for e in JsonSchemaEnum] + + if v not in valid_types: + raise ValueError(f"Invalid type: '{v}'. Must be one of {valid_types}") + return v + def simple_input_dump(self) -> dict: """Serializes the attribute into a standard JSON Schema dictionary.""" res: dict = {"type": self.type} @@ -60,17 +59,28 @@ ToolAttr.model_rebuild() class ToolCall(BaseModel): - """Model representing a tool definition and its call structure.""" + """ + Model representing a tool definition and its call structure. + Supports parsing from standard JSON Schema formats and converting to MCP Tool objects. + """ index: int = 0 id: str = "" type: str = "function" name: str = "" - arguments: str = Field(default="", description="JSON string of tool execution arguments") description: str = "" - input_schema: Dict[str, ToolAttr] = Field(default_factory=dict) - input_required: List[str] = Field(default_factory=list) - output_schema: Dict[str, ToolAttr] = Field(default_factory=dict) + + arguments: str = Field(default="", description="JSON string of tool execution arguments") + + parameters: ToolAttr = Field( + default_factory=lambda: ToolAttr(type="object", properties={}, required=[]), + description="Specification for input parameters", + ) + + output: ToolAttr = Field( + default_factory=lambda: ToolAttr(type="object", properties={}), + description="Specification for the execution result (Schema)", + ) @model_validator(mode="before") @classmethod @@ -80,25 +90,24 @@ class ToolCall(BaseModel): t_type = data.get("type", "function") body = data.get(t_type, {}) + # Extract basic metadata data["name"] = body.get("name", data.get("name", "")) data["arguments"] = body.get("arguments", data.get("arguments", "")) data["description"] = body.get("description", data.get("description", "")) + # Handle parameters mapping if "parameters" in body: params = body["parameters"] - data["input_required"] = params.get("required", []) - data["input_schema"] = {k: ToolAttr(**v) for k, v in params.get("properties", {}).items()} + # If parameters is already a dict, ensure it matches ToolAttr structure + if isinstance(params, dict): + data["parameters"] = ToolAttr(**params) + + # Handle output mapping (if provided in source) + if "output" in body and isinstance(body["output"], dict): + data["output"] = ToolAttr(**body["output"]) return data - def _build_full_schema(self) -> dict: - """Generates the top-level JSON Schema object for tool parameters.""" - return { - "type": "object", - "properties": {k: v.simple_input_dump() for k, v in self.input_schema.items()}, - "required": self.input_required, - } - def simple_input_dump(self) -> dict: """Returns a standardized tool definition dictionary.""" return { @@ -106,19 +115,18 @@ class ToolCall(BaseModel): self.type: { "name": self.name, "description": self.description, - "parameters": self._build_full_schema(), + "parameters": self.parameters.simple_input_dump(), }, } @classmethod def from_mcp_tool(cls, tool: Tool) -> "ToolCall": """Creates a ToolCall instance from an MCP Tool object.""" - schema = tool.inputSchema + # MCP Tool inputSchema maps directly to our parameters ToolAttr return cls( name=tool.name, description=tool.description or "", - input_schema={k: ToolAttr(**v) for k, v in schema.get("properties", {}).items()}, - input_required=schema.get("required", []), + parameters=ToolAttr(**tool.inputSchema), ) def to_mcp_tool(self) -> Tool: @@ -126,7 +134,7 @@ class ToolCall(BaseModel): return Tool( name=self.name, description=self.description, - inputSchema=self._build_full_schema(), + inputSchema=self.parameters.simple_input_dump(), ) @property diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index 62a9d5a0..ed1916b2 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -1,6 +1,14 @@ """utils""" -from .timer import timer +from .case_converter import snake_to_camel, camel_to_snake +from .env_utils import load_env from .singleton import singleton +from .timer import timer -__all__ = ["timer", "singleton"] +__all__ = [ + "snake_to_camel", + "camel_to_snake", + "load_env", + "singleton", + "timer", +] diff --git a/reme_ai/core/utils/env_utils.py b/reme_ai/core/utils/env_utils.py new file mode 100644 index 00000000..433039cd --- /dev/null +++ b/reme_ai/core/utils/env_utils.py @@ -0,0 +1,63 @@ +"""Environment variable loader utility for managing .env files.""" + +import os +from pathlib import Path + +from loguru import logger + +# Global flag to ensure environment is loaded only once +_ENV_LOADED = False + + +def _parse_env_file(path: Path) -> None: + """Parse and inject key-value pairs from a .env file into os.environ.""" + try: + with path.open(encoding="utf-8") as file: + for line in file: + line = line.strip() + if not line or line.startswith("#"): + continue + + if "=" in line: + key, value = line.split("=", 1) + # Strip whitespace and common quotes + os.environ[key.strip()] = value.strip().strip("'\"") + except PermissionError as err: + logger.warning(f"Permission denied for {path}: {err}") + except Exception as err: + logger.error(f"Failed to load {path}: {err}") + raise + + +def load_env(path: str | Path | None = None, enable_log: bool = True) -> None: + """Search and load the .env file into the system environment.""" + global _ENV_LOADED # pylint: disable=global-statement + if _ENV_LOADED: + return + + if path: + path = Path(path) + if path.exists(): + _parse_env_file(path) + _ENV_LOADED = True + else: + logger.warning(f".env not found at: {path}") + return + + # Search current directory and up to 5 levels of parents + for directory in [Path.cwd(), *Path.cwd().parents[:5]]: + env_path = directory / ".env" + if env_path.exists(): + if enable_log: + logger.info(f"Loading environment from: {env_path}") + _parse_env_file(env_path) + _ENV_LOADED = True + return + + logger.warning(".env file not found in search path") + + +def reset_env_flag() -> None: + """Reset the internal load state flag.""" + global _ENV_LOADED # pylint: disable=global-statement + _ENV_LOADED = False diff --git a/tests/test_llm.py b/tests/test_llm.py new file mode 100644 index 00000000..308a69a9 --- /dev/null +++ b/tests/test_llm.py @@ -0,0 +1,422 @@ +""" +Async unit tests for LLM classes (OpenAILLM and LiteLLM) covering: +- Async non-streaming chat +- Async chat with stream print +- Async streaming chat +- Async chat with tools + +Usage: + python test_llm.py --openai # Test OpenAILLM only + python test_llm.py --litellm # Test LiteLLM only + python test_llm.py --all # Test both LLMs +""" + +# flake8: noqa: E402 +# pylint: disable=C0413 + +import asyncio +import argparse +from typing import Type + +from reme_ai.core.utils import load_env + +load_env() + +from reme_ai.core.llm import OpenAILLM, LiteLLM, BaseLLM +from reme_ai.core.schema import Message, ToolCall +from reme_ai.core.enumeration import Role, ChunkEnum + + +def get_llm(llm_class: Type[BaseLLM]) -> BaseLLM: + """Create and return an LLM instance.""" + return llm_class( + model_name="qwen3-30b-a3b-instruct-2507", + max_retries=2, + raise_exception=True, + ) + + +def get_multi_turn_messages() -> list[Message]: + """Create multi-turn conversation messages for testing.""" + return [ + Message( + role=Role.SYSTEM, + content="You are a helpful AI assistant with expertise in mathematics, science, and general knowledge.", + ), + Message( + role=Role.USER, + content="Hello! I'm working on a science project about renewable energy. " + "Can you help me understand the basics?", + ), + Message( + role=Role.ASSISTANT, + content="Of course! I'd be happy to help. Renewable energy comes from sources that naturally replenish, " + "like solar, wind, hydro, geothermal, and biomass. What specific aspect would you like to explore?", + ), + Message( + role=Role.USER, + content="I'm particularly interested in solar energy. Can you explain how solar panels work and calculate " + "how much energy a typical home solar system might produce?", + ), + Message( + role=Role.ASSISTANT, + content="Solar panels work through photovoltaic cells that convert sunlight into electricity. " + "When photons hit the silicon cells, they knock electrons loose, creating an electric current." + "\n\nFor energy calculation: A typical home solar system is 5-10kW. With average 4-5 peak sun " + "hours per day, a 6kW system would produce approximately 24-30 kWh daily, or 720-900 kWh monthly.", + ), + Message( + role=Role.USER, + content="That's helpful! Now, given that calculation, if electricity costs $0.12 per kWh, " + "estimate the annual savings. Also, briefly mention what factors might affect this.", + ), + ] + + +def get_test_tools() -> list[ToolCall]: + """Create comprehensive test tools for tool calling.""" + return [ + ToolCall( + **{ + "type": "function", + "function": { + "name": "calculate_energy_savings", + "description": "Calculate annual energy savings based on solar production and electricity rates", + "parameters": { + "type": "object", + "properties": { + "monthly_kwh": { + "type": "number", + "description": "Monthly energy production in kWh", + }, + "electricity_rate": { + "type": "number", + "description": "Electricity cost per kWh in dollars", + }, + "system_efficiency": { + "type": "number", + "description": "System efficiency factor (0-1), defaults to 0.85", + }, + }, + "required": ["monthly_kwh", "electricity_rate"], + }, + }, + }, + ), + ToolCall( + **{ + "type": "function", + "function": { + "name": "get_weather_data", + "description": "Get current weather and solar irradiance data for a location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City name or coordinates, e.g., 'San Francisco' or '37.7749,-122.4194'", + }, + "include_forecast": { + "type": "boolean", + "description": "Whether to include 7-day forecast", + }, + "unit": { + "type": "string", + "description": "Temperature unit", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["location"], + }, + }, + }, + ), + ToolCall( + **{ + "type": "function", + "function": { + "name": "analyze_panel_efficiency", + "description": "Analyze solar panel efficiency based on various environmental factors", + "parameters": { + "type": "object", + "properties": { + "panel_type": { + "type": "string", + "description": "Type of solar panel", + "enum": ["monocrystalline", "polycrystalline", "thin-film"], + }, + "temperature": { + "type": "number", + "description": "Ambient temperature in Celsius", + }, + "age_years": { + "type": "number", + "description": "Age of the panel in years", + }, + }, + "required": ["panel_type", "temperature"], + }, + }, + }, + ), + ] + + +def get_tool_test_messages() -> list[Message]: + """Create multi-turn messages that should trigger tool calling.""" + return [ + Message( + role=Role.SYSTEM, + content="You are a helpful assistant with access to weather and energy calculation tools. Use them when " + "appropriate.", + ), + Message( + role=Role.USER, + content="I'm planning to install solar panels in San Francisco. Can you help me understand the weather " + "patterns there?", + ), + Message( + role=Role.ASSISTANT, + content="I'd be happy to help! San Francisco has a Mediterranean climate with " + "mild temperatures year-round. Let me get the current weather data for you.", + ), + Message( + role=Role.USER, + content="Great! Also, I'm considering monocrystalline panels. If my system produces 800 kWh monthly " + "and electricity costs $0.15 per kWh, what would be my annual savings?", + ), + ] + + +async def test_async_chat(llm_class: Type[BaseLLM], llm_name: str): + """Test asynchronous non-streaming chat with multi-turn conversation.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Async Non-Streaming Chat") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_multi_turn_messages() + + print(f"Input: {len(messages)} messages in conversation") + print(f"Last user message: {messages[-1].content[:100]}...") + + response = await llm.chat(messages=messages) + + assert response is not None, f"{llm_name}: Response is None" + assert response.role == Role.ASSISTANT, f"{llm_name}: Wrong role" + assert isinstance(response.content, str), f"{llm_name}: Content is not string" + assert len(response.content) > 0, f"{llm_name}: Empty response" + + print(f"\nResponse preview: {response.content[:200]}...") + print(f"Full response length: {len(response.content)} characters") + print(f"\nFull message:\n{response.simple_dump()}") + + await llm.close() + print(f"✓ PASSED: {llm_name} async chat") + + +async def test_async_chat_with_stream_print(llm_class: Type[BaseLLM], llm_name: str): + """Test asynchronous chat with stream print enabled.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Async Chat with Stream Print") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_multi_turn_messages() + + print(f"Input: {len(messages)} messages in conversation") + print(f"Last user message: {messages[-1].content[:100]}...") + print("\nStreaming output:") + print("-" * 60) + + response = await llm.chat(messages=messages, enable_stream_print=True) + + print("\n" + "-" * 60) + + assert response is not None, f"{llm_name}: Response is None" + assert response.role == Role.ASSISTANT, f"{llm_name}: Wrong role" + assert isinstance(response.content, str), f"{llm_name}: Content is not string" + assert len(response.content) > 0, f"{llm_name}: Empty response" + + print(f"\nFull message:\n{response.simple_dump()}") + + await llm.close() + print(f"✓ PASSED: {llm_name} async chat with stream print") + + +async def test_async_stream_chat(llm_class: Type[BaseLLM], llm_name: str): + """Test asynchronous streaming chat.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Async Streaming Chat") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_multi_turn_messages() + + print(f"Input: {len(messages)} messages in conversation") + print(f"Last user message: {messages[-1].content[:100]}...") + print("\nStreaming chunks:") + print("-" * 60) + + chunks = [] + answer_content = "" + + async for chunk in llm.stream_chat(messages=messages): + chunks.append(chunk) + if chunk.chunk_type == ChunkEnum.ANSWER: + answer_content += chunk.chunk + print(chunk.chunk, end="", flush=True) + + print("\n" + "-" * 60) + + assert len(chunks) > 0, f"{llm_name}: No chunks received" + assert len(answer_content) > 0, f"{llm_name}: Empty answer content" + + # Check that we received at least one ANSWER or USAGE chunk + chunk_types = [c.chunk_type for c in chunks] + assert ChunkEnum.ANSWER in chunk_types or ChunkEnum.USAGE in chunk_types, f"{llm_name}: No ANSWER or USAGE chunks" + + # Print the final assembled message + if chunks and hasattr(chunks[-1], "message") and chunks[-1].message: + print(f"\nFull message:\n{chunks[-1].message.simple_dump()}") + + print(f"\nTotal chunks: {len(chunks)}") + print(f"Answer length: {len(answer_content)} characters") + + await llm.close() + print(f"✓ PASSED: {llm_name} async streaming chat") + + +async def test_async_chat_with_tools(llm_class: Type[BaseLLM], llm_name: str): + """Test asynchronous chat with tool calling.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Async Chat with Tools") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_tool_test_messages() + tools = get_test_tools() + + print(f"Input: {len(messages)} messages, {len(tools)} tools available") + print(f"Tools: {[tool.name for tool in tools]}") + print(f"Last user message: {messages[-1].content[:100]}...") + + response = await llm.chat(messages=messages, tools=tools) + + assert response is not None, f"{llm_name}: Response is None" + assert response.role == Role.ASSISTANT, f"{llm_name}: Wrong role" + # Response should contain either content or tool_calls + assert response.content or response.tool_calls, f"{llm_name}: No content or tool_calls" + + if response.tool_calls: + print(f"\n✓ Tool calls detected: {len(response.tool_calls)}") + for i, tool_call in enumerate(response.tool_calls, 1): + print(f"\n Tool call #{i}:") + print(f" - Name: {tool_call.name}") + print(f" - Arguments: {tool_call.arguments}") + # Validate that arguments are valid JSON + assert tool_call.check_argument(), f"{llm_name}: Invalid tool arguments" + print(" - ✓ Arguments validated") + else: + print("\n⚠ No tool calls (response with text instead)") + print(f"Response preview: {response.content[:200]}...") + + print(f"\nFull message:\n{response.simple_dump()}") + + await llm.close() + print(f"✓ PASSED: {llm_name} async chat with tools") + + +async def run_all_tests_for_llm(llm_class: Type[BaseLLM], llm_name: str): + """Run all tests for a specific LLM class.""" + print(f"\n\n{'#'*60}") + print(f"# Running all tests for: {llm_name}") + print(f"{'#'*60}") + + await test_async_chat(llm_class, llm_name) + await test_async_chat_with_stream_print(llm_class, llm_name) + await test_async_stream_chat(llm_class, llm_name) + await test_async_chat_with_tools(llm_class, llm_name) + + print(f"\n{'='*60}") + print(f"✓ All tests passed for {llm_name}!") + print(f"{'='*60}") + + +async def main(): + """Main entry point for running tests.""" + parser = argparse.ArgumentParser( + description="Run async LLM tests", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python test_llm.py --openai # Test OpenAILLM only + python test_llm.py --litellm # Test LiteLLM only + python test_llm.py --all # Test both LLMs + """, + ) + parser.add_argument( + "--openai", + action="store_true", + help="Test OpenAILLM", + ) + parser.add_argument( + "--litellm", + action="store_true", + help="Test LiteLLM", + ) + parser.add_argument( + "--all", + action="store_true", + help="Run tests for all available LLMs", + ) + + args = parser.parse_args() + + # Determine which LLMs to test + llms_to_test = [] + + if args.all: + llms_to_test = [ + (OpenAILLM, "OpenAILLM"), + (LiteLLM, "LiteLLM"), + ] + elif args.openai and args.litellm: + llms_to_test = [ + (OpenAILLM, "OpenAILLM"), + (LiteLLM, "LiteLLM"), + ] + elif args.openai: + llms_to_test = [(OpenAILLM, "OpenAILLM")] + elif args.litellm: + llms_to_test = [(LiteLLM, "LiteLLM")] + else: + # Default to all LLMs if no argument provided + llms_to_test = [ + (OpenAILLM, "OpenAILLM"), + (LiteLLM, "LiteLLM"), + ] + print("No LLM specified, defaulting to --all (testing all LLMs)") + print("Use --openai or --litellm to test a specific one\n") + + # Run tests for each LLM + for llm_class, llm_name in llms_to_test: + try: + await run_all_tests_for_llm(llm_class, llm_name) + except Exception as e: + print(f"\n✗ FAILED: {llm_name} tests failed with error:") + print(f" {type(e).__name__}: {e}") + raise + + # Final summary + print(f"\n\n{'#'*60}") + print("# TEST SUMMARY") + print(f"{'#'*60}") + print(f"✓ All tests passed for {len(llms_to_test)} LLM(s):") + for _, llm_name in llms_to_test: + print(f" - {llm_name}") + print(f"{'#'*60}\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_llm_sync.py b/tests/test_llm_sync.py new file mode 100644 index 00000000..98751f87 --- /dev/null +++ b/tests/test_llm_sync.py @@ -0,0 +1,421 @@ +""" +Sync unit tests for LLM classes (OpenAILLM and LiteLLM) covering: +- Sync non-streaming chat +- Sync chat with stream print +- Sync streaming chat +- Sync chat with tools + +Usage: + python test_llm_sync.py --openai # Test OpenAILLM only + python test_llm_sync.py --litellm # Test LiteLLM only + python test_llm_sync.py --all # Test both LLMs +""" + +# flake8: noqa: E402 +# pylint: disable=C0413 + +import argparse +from typing import Type + +from reme_ai.core.utils import load_env + +load_env() + +from reme_ai.core.llm import OpenAILLMSync, LiteLLMSync, BaseLLM +from reme_ai.core.schema import Message, ToolCall +from reme_ai.core.enumeration import Role, ChunkEnum + + +def get_llm(llm_class: Type[BaseLLM]) -> BaseLLM: + """Create and return an LLM instance.""" + return llm_class( + model_name="qwen3-30b-a3b-instruct-2507", + max_retries=2, + raise_exception=True, + ) + + +def get_multi_turn_messages() -> list[Message]: + """Create multi-turn conversation messages for testing.""" + return [ + Message( + role=Role.SYSTEM, + content="You are a helpful AI assistant with expertise in mathematics, science, and general knowledge.", + ), + Message( + role=Role.USER, + content="Hello! I'm working on a science project about renewable energy. Can you help me understand " + "the basics?", + ), + Message( + role=Role.ASSISTANT, + content="Of course! I'd be happy to help. Renewable energy comes from sources that naturally replenish, " + "like solar, wind, hydro, geothermal, and biomass. What specific aspect would you like to explore?", + ), + Message( + role=Role.USER, + content="I'm particularly interested in solar energy. Can you explain how solar panels work " + "and calculate how much energy a typical home solar system might produce?", + ), + Message( + role=Role.ASSISTANT, + content="Solar panels work through photovoltaic cells that convert sunlight into electricity. " + "When photons hit the silicon cells, they knock electrons loose, creating an electric current." + "\n\nFor energy calculation: A typical home solar system is 5-10kW. With average 4-5 peak sun " + "hours per day, a 6kW system would produce approximately 24-30 kWh daily, or 720-900 kWh monthly.", + ), + Message( + role=Role.USER, + content="That's helpful! Now, given that calculation, if electricity costs $0.12 per kWh, " + "estimate the annual savings. Also, briefly mention what factors might affect this.", + ), + ] + + +def get_test_tools() -> list[ToolCall]: + """Create comprehensive test tools for tool calling.""" + return [ + ToolCall( + **{ + "type": "function", + "function": { + "name": "calculate_energy_savings", + "description": "Calculate annual energy savings based on solar production and electricity rates", + "parameters": { + "type": "object", + "properties": { + "monthly_kwh": { + "type": "number", + "description": "Monthly energy production in kWh", + }, + "electricity_rate": { + "type": "number", + "description": "Electricity cost per kWh in dollars", + }, + "system_efficiency": { + "type": "number", + "description": "System efficiency factor (0-1), defaults to 0.85", + }, + }, + "required": ["monthly_kwh", "electricity_rate"], + }, + }, + }, + ), + ToolCall( + **{ + "type": "function", + "function": { + "name": "get_weather_data", + "description": "Get current weather and solar irradiance data for a location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City name or coordinates, e.g., 'San Francisco' or '37.7749,-122.4194'", + }, + "include_forecast": { + "type": "boolean", + "description": "Whether to include 7-day forecast", + }, + "unit": { + "type": "string", + "description": "Temperature unit", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["location"], + }, + }, + }, + ), + ToolCall( + **{ + "type": "function", + "function": { + "name": "analyze_panel_efficiency", + "description": "Analyze solar panel efficiency based on various environmental factors", + "parameters": { + "type": "object", + "properties": { + "panel_type": { + "type": "string", + "description": "Type of solar panel", + "enum": ["monocrystalline", "polycrystalline", "thin-film"], + }, + "temperature": { + "type": "number", + "description": "Ambient temperature in Celsius", + }, + "age_years": { + "type": "number", + "description": "Age of the panel in years", + }, + }, + "required": ["panel_type", "temperature"], + }, + }, + }, + ), + ] + + +def get_tool_test_messages() -> list[Message]: + """Create multi-turn messages that should trigger tool calling.""" + return [ + Message( + role=Role.SYSTEM, + content="You are a helpful assistant with access to weather and energy calculation tools. " + "Use them when appropriate.", + ), + Message( + role=Role.USER, + content="I'm planning to install solar panels in San Francisco. Can you help me understand the " + "weather patterns there?", + ), + Message( + role=Role.ASSISTANT, + content="I'd be happy to help! San Francisco has a Mediterranean climate with mild temperatures " + "year-round. Let me get the current weather data for you.", + ), + Message( + role=Role.USER, + content="Great! Also, I'm considering monocrystalline panels. If my system produces 800 kWh monthly " + "and electricity costs $0.15 per kWh, what would be my annual savings?", + ), + ] + + +def test_sync_chat(llm_class: Type[BaseLLM], llm_name: str): + """Test synchronous non-streaming chat with multi-turn conversation.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Sync Non-Streaming Chat") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_multi_turn_messages() + + print(f"Input: {len(messages)} messages in conversation") + print(f"Last user message: {messages[-1].content[:100]}...") + + response = llm.chat_sync(messages=messages) + + assert response is not None, f"{llm_name}: Response is None" + assert response.role == Role.ASSISTANT, f"{llm_name}: Wrong role" + assert isinstance(response.content, str), f"{llm_name}: Content is not string" + assert len(response.content) > 0, f"{llm_name}: Empty response" + + print(f"\nResponse preview: {response.content[:200]}...") + print(f"Full response length: {len(response.content)} characters") + print(f"\nFull message:\n{response.simple_dump()}") + + llm.close_sync() + print(f"✓ PASSED: {llm_name} sync chat") + + +def test_sync_chat_with_stream_print(llm_class: Type[BaseLLM], llm_name: str): + """Test synchronous chat with stream print enabled.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Sync Chat with Stream Print") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_multi_turn_messages() + + print(f"Input: {len(messages)} messages in conversation") + print(f"Last user message: {messages[-1].content[:100]}...") + print("\nStreaming output:") + print("-" * 60) + + response = llm.chat_sync(messages=messages, enable_stream_print=True) + + print("\n" + "-" * 60) + + assert response is not None, f"{llm_name}: Response is None" + assert response.role == Role.ASSISTANT, f"{llm_name}: Wrong role" + assert isinstance(response.content, str), f"{llm_name}: Content is not string" + assert len(response.content) > 0, f"{llm_name}: Empty response" + + print(f"\nFull message:\n{response.simple_dump()}") + + llm.close_sync() + print(f"✓ PASSED: {llm_name} sync chat with stream print") + + +def test_sync_stream_chat(llm_class: Type[BaseLLM], llm_name: str): + """Test synchronous streaming chat.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Sync Streaming Chat") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_multi_turn_messages() + + print(f"Input: {len(messages)} messages in conversation") + print(f"Last user message: {messages[-1].content[:100]}...") + print("\nStreaming chunks:") + print("-" * 60) + + chunks = [] + answer_content = "" + + for chunk in llm.stream_chat_sync(messages=messages): + chunks.append(chunk) + if chunk.chunk_type == ChunkEnum.ANSWER: + answer_content += chunk.chunk + print(chunk.chunk, end="", flush=True) + + print("\n" + "-" * 60) + + assert len(chunks) > 0, f"{llm_name}: No chunks received" + assert len(answer_content) > 0, f"{llm_name}: Empty answer content" + + # Check that we received at least one ANSWER or USAGE chunk + chunk_types = [c.chunk_type for c in chunks] + assert ChunkEnum.ANSWER in chunk_types or ChunkEnum.USAGE in chunk_types, f"{llm_name}: No ANSWER or USAGE chunks" + + # Print the final assembled message + if chunks and hasattr(chunks[-1], "message") and chunks[-1].message: + print(f"\nFull message:\n{chunks[-1].message.simple_dump()}") + + print(f"\nTotal chunks: {len(chunks)}") + print(f"Answer length: {len(answer_content)} characters") + + llm.close_sync() + print(f"✓ PASSED: {llm_name} sync streaming chat") + + +def test_sync_chat_with_tools(llm_class: Type[BaseLLM], llm_name: str): + """Test synchronous chat with tool calling.""" + print(f"\n{'='*60}") + print(f"Testing {llm_name}: Sync Chat with Tools") + print(f"{'='*60}") + + llm = get_llm(llm_class) + messages = get_tool_test_messages() + tools = get_test_tools() + + print(f"Input: {len(messages)} messages, {len(tools)} tools available") + print(f"Tools: {[tool.name for tool in tools]}") + print(f"Last user message: {messages[-1].content[:100]}...") + + response = llm.chat_sync(messages=messages, tools=tools) + + assert response is not None, f"{llm_name}: Response is None" + assert response.role == Role.ASSISTANT, f"{llm_name}: Wrong role" + # Response should contain either content or tool_calls + assert response.content or response.tool_calls, f"{llm_name}: No content or tool_calls" + + if response.tool_calls: + print(f"\n✓ Tool calls detected: {len(response.tool_calls)}") + for i, tool_call in enumerate(response.tool_calls, 1): + print(f"\n Tool call #{i}:") + print(f" - Name: {tool_call.name}") + print(f" - Arguments: {tool_call.arguments}") + # Validate that arguments are valid JSON + assert tool_call.check_argument(), f"{llm_name}: Invalid tool arguments" + print(" - ✓ Arguments validated") + else: + print("\n⚠ No tool calls (response with text instead)") + print(f"Response preview: {response.content[:200]}...") + + print(f"\nFull message:\n{response.simple_dump()}") + + llm.close_sync() + print(f"✓ PASSED: {llm_name} sync chat with tools") + + +def run_all_tests_for_llm(llm_class: Type[BaseLLM], llm_name: str): + """Run all tests for a specific LLM class.""" + print(f"\n\n{'#'*60}") + print(f"# Running all tests for: {llm_name}") + print(f"{'#'*60}") + + test_sync_chat(llm_class, llm_name) + test_sync_chat_with_stream_print(llm_class, llm_name) + test_sync_stream_chat(llm_class, llm_name) + test_sync_chat_with_tools(llm_class, llm_name) + + print(f"\n{'='*60}") + print(f"✓ All tests passed for {llm_name}!") + print(f"{'='*60}") + + +def main(): + """Main entry point for running tests.""" + parser = argparse.ArgumentParser( + description="Run sync LLM tests", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python test_llm_sync.py --openai # Test OpenAILLM only + python test_llm_sync.py --litellm # Test LiteLLM only + python test_llm_sync.py --all # Test both LLMs + """, + ) + parser.add_argument( + "--openai", + action="store_true", + help="Test OpenAILLM", + ) + parser.add_argument( + "--litellm", + action="store_true", + help="Test LiteLLM", + ) + parser.add_argument( + "--all", + action="store_true", + help="Run tests for all available LLMs", + ) + + args = parser.parse_args() + + # Determine which LLMs to test + llms_to_test = [] + + if args.all: + llms_to_test = [ + (OpenAILLMSync, "OpenAILLMSync"), + (LiteLLMSync, "LiteLLMSync"), + ] + elif args.openai and args.litellm: + llms_to_test = [ + (OpenAILLMSync, "OpenAILLMSync"), + (LiteLLMSync, "LiteLLMSync"), + ] + elif args.openai: + llms_to_test = [(OpenAILLMSync, "OpenAILLMSync")] + elif args.litellm: + llms_to_test = [(LiteLLMSync, "LiteLLMSync")] + else: + # Default to all LLMs if no argument provided + llms_to_test = [ + (OpenAILLMSync, "OpenAILLMSync"), + (LiteLLMSync, "LiteLLMSync"), + ] + print("No LLM specified, defaulting to --all (testing all LLMs)") + print("Use --openai or --litellm to test a specific one\n") + + # Run tests for each LLM + for llm_class, llm_name in llms_to_test: + try: + run_all_tests_for_llm(llm_class, llm_name) + except Exception as e: + print(f"\n✗ FAILED: {llm_name} tests failed with error:") + print(f" {type(e).__name__}: {e}") + raise + + # Final summary + print(f"\n\n{'#'*60}") + print("# TEST SUMMARY") + print(f"{'#'*60}") + print(f"✓ All tests passed for {len(llms_to_test)} LLM(s):") + for _, llm_name in llms_to_test: + print(f" - {llm_name}") + print(f"{'#'*60}\n") + + +if __name__ == "__main__": + main() diff --git a/tests/test_message.py b/tests/test_message.py index b45563dd..141174c5 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -75,11 +75,11 @@ class TestModelDefinitions(unittest.TestCase): print(tc.simple_output_dump()) self.assertEqual(tc.name, "get_weather") - self.assertIn("location", tc.input_schema) - self.assertIn("unit", tc.input_schema) + self.assertIn("location", tc.parameters.properties) + self.assertIn("unit", tc.parameters.properties) # Check that 'location' is in the required list at ToolCall level - self.assertIn("location", tc.input_required) - self.assertNotIn("unit", tc.input_required) + self.assertIn("location", tc.parameters.required) + self.assertNotIn("unit", tc.parameters.required) def test_tool_call_argument_parsing(self): """Test JSON argument parsing and validation.""" @@ -170,11 +170,11 @@ class TestModelDefinitions(unittest.TestCase): print(tc.simple_output_dump()) self.assertEqual(tc.name, "calculator") - self.assertIn("a", tc.input_schema) - self.assertIn("b", tc.input_schema) + self.assertIn("a", tc.parameters.properties) + self.assertIn("b", tc.parameters.properties) # Check that 'a' is in the required list - self.assertIn("a", tc.input_required) - self.assertNotIn("b", tc.input_required) + self.assertIn("a", tc.parameters.required) + self.assertNotIn("b", tc.parameters.required) # From ToolCall back to MCP structure (via to_mcp_tool) # Note: This checks the logic of constructing the dict for Tool(...) diff --git a/tests/test_tool_call.py b/tests/test_tool_call.py index f914dca7..30c0a37e 100644 --- a/tests/test_tool_call.py +++ b/tests/test_tool_call.py @@ -28,7 +28,7 @@ def test_simple_schema(): # 解析 tool_call = ToolCall.model_validate(raw_definition) print(f"工具名称: {tool_call.name}") - print(f"必填参数: {tool_call.input_required}") + print(f"必填参数: {tool_call.parameters.required}") # 导出并验证相等性 dumped_data = tool_call.simple_input_dump() @@ -73,9 +73,9 @@ def test_medium_nested_schema(): # 解析 tool_call = ToolCall.model_validate(raw_definition) print(f"工具名称: {tool_call.name}") - print(f"根级必填项: {tool_call.input_required}") + print(f"根级必填项: {tool_call.parameters.required}") - customer_attr = tool_call.input_schema["customer"] + customer_attr = tool_call.parameters.properties["customer"] print(f"Customer 子属性: {list(customer_attr.properties.keys())}") print(f"Customer 必填项: {customer_attr.required}") @@ -134,10 +134,10 @@ def test_nested_schema(): tool_call = ToolCall.model_validate(raw_definition) print(f"工具名称: {tool_call.name}") - print(f"根级必填项: {tool_call.input_required}") + print(f"根级必填项: {tool_call.parameters.required}") # 验证嵌套深度 - metadata_attr = tool_call.input_schema["metadata"] + metadata_attr = tool_call.parameters.properties["metadata"] print(f"Metadata 子属性: {list(metadata_attr.properties.keys())}") print(f"Metadata 必填项: {metadata_attr.required}") @@ -189,9 +189,9 @@ def test_array_of_primitives(): # 解析 tool_call = ToolCall.model_validate(raw_definition) print(f"工具名称: {tool_call.name}") - print(f"必填参数: {tool_call.input_required}") + print(f"必填参数: {tool_call.parameters.required}") - file_paths_attr = tool_call.input_schema["file_paths"] + file_paths_attr = tool_call.parameters.properties["file_paths"] print(f"file_paths 类型: {file_paths_attr.type}") t_items_type = file_paths_attr.items.type if hasattr(file_paths_attr.items, "type") else file_paths_attr.items print(f"file_paths items 类型: {t_items_type}") @@ -267,9 +267,9 @@ def test_deep_nested_schema(): # 解析 tool_call = ToolCall.model_validate(raw_definition) print(f"工具名称: {tool_call.name}") - print(f"根级必填项: {tool_call.input_required}") + print(f"根级必填项: {tool_call.parameters.required}") - team_attr = tool_call.input_schema["team"] + team_attr = tool_call.parameters.properties["team"] leader_attr = team_attr.properties["leader"] contact_attr = leader_attr.properties["contact"] print(f"Team 必填项: {team_attr.required}") @@ -330,13 +330,13 @@ def test_mixed_types_schema(): # 解析 tool_call = ToolCall.model_validate(raw_definition) print(f"工具名称: {tool_call.name}") - print(f"根级必填项: {tool_call.input_required}") + print(f"根级必填项: {tool_call.parameters.required}") # 验证各种类型 - print(f"enabled 类型: {tool_call.input_schema['enabled'].type}") - print(f"max_connections 类型: {tool_call.input_schema['max_connections'].type}") - print(f"timeout 类型: {tool_call.input_schema['timeout'].type}") - print(f"mode 枚举值: {tool_call.input_schema['mode'].enum}") + print(f"enabled 类型: {tool_call.parameters.properties['enabled'].type}") + print(f"max_connections 类型: {tool_call.parameters.properties['max_connections'].type}") + print(f"timeout 类型: {tool_call.parameters.properties['timeout'].type}") + print(f"mode 枚举值: {tool_call.parameters.properties['mode'].enum}") # 导出并验证相等性 dumped_data = tool_call.simple_input_dump() From c49665cbd41971184bb523bb9f7b41ef6d0b7022 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 00:54:40 +0800 Subject: [PATCH 07/10] docs(schema): update tool call documentation with example --- reme_ai/core/schema/tool_call.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/reme_ai/core/schema/tool_call.py b/reme_ai/core/schema/tool_call.py index e6b1092f..a94dfb1d 100644 --- a/reme_ai/core/schema/tool_call.py +++ b/reme_ai/core/schema/tool_call.py @@ -62,6 +62,34 @@ class ToolCall(BaseModel): """ Model representing a tool definition and its call structure. Supports parsing from standard JSON Schema formats and converting to MCP Tool objects. + input: + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "It is very useful when you want to check the weather of a specified city.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.", + } + }, + "required": ["location"] + } + } + } + output: + { + "index": 0, + "id": "call_6596dafa2a6a46f7a217da", + "function": { + "arguments": "{\"location\": \"Beijing\"}", + "name": "get_current_weather" + }, + "type": "function", + } """ index: int = 0 From 5f6244d43265dedd22afa659470726a41d630de0 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 00:55:22 +0800 Subject: [PATCH 08/10] ci(workflow): add pre-commit workflow for code formatting --- .github/workflows/pre-commit.yml | 38 ++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 .github/workflows/pre-commit.yml diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml new file mode 100644 index 00000000..9685f577 --- /dev/null +++ b/.github/workflows/pre-commit.yml @@ -0,0 +1,38 @@ +name: Pre-commit + +on: [ push, pull_request ] + +jobs: + run: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: True + matrix: + os: [ ubuntu-latest ] + env: + OS: ${{ matrix.os }} + PYTHON: '3.10' + steps: + - uses: actions/checkout@master + - name: Setup Python + uses: actions/setup-python@master + with: + python-version: '3.10' + - name: Update setuptools + run: | + pip install -U setuptools wheel + - name: Install + run: | + pip install -q -e .[dev] + - name: Install pre-commit + run: | + pre-commit install + - name: Pre-commit starts + run: | + pre-commit run --all-files > pre-commit.log 2>&1 || true + cat pre-commit.log + if grep -q Failed pre-commit.log; then + echo -e "\e[41m [**FAIL**] Please install pre-commit and format your code first. \e[0m" + exit 1 + fi + echo -e "\e[46m ********************************Passed******************************** \e[0m" \ No newline at end of file From c655f366e3593e5e96dd44dc7e2d47be23d87da5 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 09:44:45 +0800 Subject: [PATCH 09/10] docs(readme): add download count and commit activity badges --- README.md | 2 ++ README_ZH.md | 2 ++ 2 files changed, 4 insertions(+) diff --git a/README.md b/README.md index 2becbcff..2423c643 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,8 @@

Python Version PyPI Version + PyPI Downloads + GitHub commit activity License English 简体中文 diff --git a/README_ZH.md b/README_ZH.md index c59b7186..5a39d3e6 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -5,6 +5,8 @@

Python 版本 PyPI 版本 + PyPI Downloads + GitHub commit activity 许可证 English 简体中文 From 38c972b270d36c42a7a7fb265d384be5507ff038 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 31 Dec 2025 09:48:13 +0800 Subject: [PATCH 10/10] chore(deps): add pre-commit to project dependencies --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 60de58d9..abb263da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,8 @@ dev = [ "myst-nb", "sphinxcontrib-bibtex", "furo", - "sphinxcontrib-mermaid" + "sphinxcontrib-mermaid", + "pre-commit", ] token = [