scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-09-12 23:01:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor(llm): update type hints and abstract method definitions

Browse source
This commit is contained in:
jinli.yl 2026-01-08 11:37:02 +08:00
parent 554eec1cb9
commit b3d68dbc02
5 changed files with 88 additions and 156 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
reme_ai/core/llm/base_llm.py
View file

@ -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:

56
reme_ai/core/llm/lite_llm.py
View file

@ -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)

36
reme_ai/core/llm/lite_llm_sync.py
View file

@ -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)

48
reme_ai/core/llm/openai_llm.py
View file

@ -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):

42
reme_ai/core/llm/openai_llm_sync.py
View file

@ -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):