diff --git a/reme_ai/core/llm/base_llm.py b/reme_ai/core/llm/base_llm.py index a7ba0f76..5235198c 100644 --- a/reme_ai/core/llm/base_llm.py +++ b/reme_ai/core/llm/base_llm.py @@ -3,8 +3,8 @@ import asyncio import json import time -from abc import ABC -from typing import Callable, Generator, AsyncGenerator, Optional, Any +from abc import ABC, abstractmethod +from typing import Callable, Generator, AsyncGenerator, Any from loguru import logger @@ -72,10 +72,7 @@ class BaseLLM(ABC): ) @staticmethod - def _accumulate_tool_call_chunk( - tool_call, - ret_tools: list[ToolCall], - ) -> None: + def _accumulate_tool_call_chunk(tool_call, ret_tools: list[ToolCall]): """Assemble incremental tool call fragments into complete ToolCall objects.""" index = tool_call.index @@ -94,48 +91,39 @@ class BaseLLM(ABC): ret_tools[index].arguments += tool_call.function.arguments @staticmethod - def _validate_and_serialize_tools( - ret_tools: list[ToolCall], - tools: Optional[list[ToolCall]], - ) -> list[dict]: + def _validate_and_serialize_tools(ret_tool_calls: list[ToolCall], tools: list[ToolCall]) -> list[dict]: """Validate tool call integrity and return serialized tool dictionaries.""" - if not ret_tools: + if not ret_tool_calls: 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 + for tool in ret_tool_calls: 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}", - ) + raise ValueError(f"Tool call {tool.name} has invalid JSON arguments: {tool.arguments}") validated_tools.append(tool.simple_output_dump()) - return validated_tools + @abstractmethod def _build_stream_kwargs( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = 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, + tools: list[ToolCall] | None, + stream_kwargs: dict, ) -> AsyncGenerator[StreamChunk, None]: """Internal async generator for streaming raw response chunks.""" raise NotImplementedError @@ -143,8 +131,8 @@ class BaseLLM(ABC): def _stream_chat_sync( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, - stream_kwargs: Optional[dict] = None, + tools: list[ToolCall] | None = None, + stream_kwargs: dict | None = None, ) -> Generator[StreamChunk, None, None]: """Internal synchronous generator for streaming raw response chunks.""" raise NotImplementedError @@ -153,7 +141,7 @@ class BaseLLM(ABC): self, operation_name: str, messages: list[Message], - tools: Optional[list[ToolCall]], + tools: list[ToolCall] | None, stream_kwargs: dict, ) -> AsyncGenerator[StreamChunk, None]: """Execute the async streaming operation with retry logic and error recovery.""" @@ -179,7 +167,7 @@ class BaseLLM(ABC): self, operation_name: str, messages: list[Message], - tools: Optional[list[ToolCall]], + tools: list[ToolCall] | None, stream_kwargs: dict, ) -> Generator[StreamChunk, None, None]: """Execute the synchronous streaming operation with retry logic and error recovery.""" @@ -203,7 +191,7 @@ class BaseLLM(ABC): async def stream_chat( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, **kwargs, ) -> AsyncGenerator[StreamChunk, None]: """Public async interface for streaming chat completions with retries.""" @@ -214,7 +202,7 @@ class BaseLLM(ABC): def stream_chat_sync( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, **kwargs, ) -> Generator[StreamChunk, None, None]: """Public synchronous interface for streaming chat completions with retries.""" @@ -224,7 +212,7 @@ class BaseLLM(ABC): async def _chat( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, enable_stream_print: bool = False, **kwargs, ) -> Message: @@ -246,7 +234,7 @@ class BaseLLM(ABC): def _chat_sync( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, enable_stream_print: bool = False, **kwargs, ) -> Message: @@ -269,7 +257,7 @@ class BaseLLM(ABC): self, operation_name: str, operation_fn: Callable[[], Any], - callback_fn: Optional[Callable[[Message], Any]] = None, + callback_fn: Callable[[Message], Any] | None = None, default_value: Any = None, ) -> Message | Any: """Execute a generic async operation with error handling and retry logic.""" @@ -293,7 +281,7 @@ class BaseLLM(ABC): self, operation_name: str, operation_fn: Callable[[], Message], - callback_fn: Optional[Callable[[Message], Any]] = None, + callback_fn: Callable[[Message], Any] | None = None, default_value: Any = None, ) -> Message | Any: """Execute a generic synchronous operation with error handling and retry logic.""" @@ -316,9 +304,9 @@ class BaseLLM(ABC): async def chat( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, enable_stream_print: bool = False, - callback_fn: Optional[Callable[[Message], Any]] = None, + callback_fn: Callable[[Message], Any] | None = None, default_value: Any = None, **kwargs, ) -> Message | Any: @@ -338,9 +326,9 @@ class BaseLLM(ABC): def chat_sync( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, enable_stream_print: bool = False, - callback_fn: Optional[Callable[[Message], Any]] = None, + callback_fn: Callable[[Message], Any] | None = None, default_value: Any = None, **kwargs, ) -> Message | Any: diff --git a/reme_ai/core/llm/lite_llm.py b/reme_ai/core/llm/lite_llm.py index 32d3005c..a04b91f3 100644 --- a/reme_ai/core/llm/lite_llm.py +++ b/reme_ai/core/llm/lite_llm.py @@ -1,7 +1,7 @@ """LiteLLM asynchronous implementation for ReMe.""" import os -from typing import List, AsyncGenerator, Optional +from typing import AsyncGenerator import litellm from loguru import logger @@ -20,21 +20,21 @@ class LiteLLM(BaseLLM): def __init__( self, - api_key: Optional[str] = None, - base_url: Optional[str] = None, + api_key: str | None = None, + base_url: str | None = 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.api_key: str | None = api_key or os.getenv("REME_LLM_API_KEY") + self.base_url: str | None = 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, + messages: list[Message], + tools: list[ToolCall] | None = None, log_params: bool = True, **kwargs, ) -> dict: @@ -73,46 +73,32 @@ class LiteLLM(BaseLLM): async def _stream_chat( self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - stream_kwargs: Optional[dict] = None, + messages: list[Message], + tools: list[ToolCall] | None = None, + stream_kwargs: dict | None = 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 + ret_tool_calls: list[ToolCall] = [] 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()) + continue - else: - delta = chunk.choices[0].delta + 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) + 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 + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - # Yield regular text content - if delta.content is not None: - yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + 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_tool_calls) - # 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): + for tool_data in self._validate_and_serialize_tools(ret_tool_calls, 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 index 8015a275..c3a925d2 100644 --- a/reme_ai/core/llm/lite_llm_sync.py +++ b/reme_ai/core/llm/lite_llm_sync.py @@ -23,41 +23,27 @@ class LiteLLMSync(LiteLLM): stream_kwargs: dict | None = None, ) -> Generator[StreamChunk, None, None]: """Internal synchronous generator for processing streaming chat completion chunks.""" - # 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 + ret_tool_calls: list[ToolCall] = [] 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()) + continue - else: - delta = chunk.choices[0].delta + 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) + 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 + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - # Yield regular text content - if delta.content is not None: - yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + 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_tool_calls) - # 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): + for tool_data in self._validate_and_serialize_tools(ret_tool_calls, 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 index df368021..ebae540c 100644 --- a/reme_ai/core/llm/openai_llm.py +++ b/reme_ai/core/llm/openai_llm.py @@ -1,7 +1,7 @@ """Asynchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content.""" import os -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator from loguru import logger from openai import AsyncOpenAI @@ -20,8 +20,8 @@ class OpenAILLM(BaseLLM): def __init__( self, - api_key: Optional[str] = None, - base_url: Optional[str] = None, + api_key: str | None = None, + base_url: str | None = None, **kwargs, ): """Initialize the OpenAI async client with API credentials and model configuration.""" @@ -39,7 +39,7 @@ class OpenAILLM(BaseLLM): def _build_stream_kwargs( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, + tools: list[ToolCall] | None = None, log_params: bool = True, **kwargs, ) -> dict: @@ -69,47 +69,33 @@ class OpenAILLM(BaseLLM): async def _stream_chat( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, - stream_kwargs: Optional[dict] = None, + tools: list[ToolCall] | None, + stream_kwargs: dict, ) -> 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 + ret_tool_calls: list[ToolCall] = [] 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()) + continue - else: - delta = chunk.choices[0].delta + 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) + 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 + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - # Yield regular text content - if delta.content is not None: - yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + if delta.tool_calls is not None: + for tool_call in delta.tool_calls: + self._accumulate_tool_call_chunk(tool_call, ret_tool_calls) - # 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): + for tool_data in self._validate_and_serialize_tools(ret_tool_calls, tools): yield StreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool_data) async def close(self): diff --git a/reme_ai/core/llm/openai_llm_sync.py b/reme_ai/core/llm/openai_llm_sync.py index 408ce272..51da29f4 100644 --- a/reme_ai/core/llm/openai_llm_sync.py +++ b/reme_ai/core/llm/openai_llm_sync.py @@ -1,6 +1,6 @@ """Synchronous OpenAI-compatible LLM implementation supporting streaming, tool calls, and reasoning content.""" -from typing import Generator, Optional +from typing import Generator from openai import OpenAI @@ -23,47 +23,33 @@ class OpenAILLMSync(OpenAILLM): def _stream_chat_sync( self, messages: list[Message], - tools: Optional[list[ToolCall]] = None, - stream_kwargs: Optional[dict] = None, + tools: list[ToolCall] | None = None, + stream_kwargs: dict | None = 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 + ret_tool_calls: list[ToolCall] = [] 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()) + continue - else: - delta = chunk.choices[0].delta + 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) + 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 + if delta.content is not None: + yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - # Yield regular text content - if delta.content is not None: - yield StreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) + if delta.tool_calls is not None: + for tool_call in delta.tool_calls: + self._accumulate_tool_call_chunk(tool_call, ret_tool_calls) - # 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): + for tool_data in self._validate_and_serialize_tools(ret_tool_calls, tools): yield StreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool_data) def close_sync(self):