mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-19 00:01:33 +00:00
510 lines
19 KiB
Python
510 lines
19 KiB
Python
"""Base interface for LLM implementations."""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from typing import Callable, Generator, AsyncGenerator, Any
|
|
|
|
from loguru import logger
|
|
|
|
from ..enumeration import ChunkEnum, Role
|
|
from ..schema import Message, StreamChunk, ToolCall
|
|
from ..utils import extract_content
|
|
|
|
|
|
class BaseLLM(ABC):
|
|
"""Base class for LLM interactions."""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
model_name: str = "",
|
|
max_retries: int = 10,
|
|
raise_exception: bool = False,
|
|
request_interval: float = 0.0,
|
|
**kwargs,
|
|
):
|
|
"""Initialize LLM client.
|
|
|
|
Args:
|
|
model_name: Model name to use
|
|
max_retries: Maximum retry attempts on failure
|
|
raise_exception: Raise exceptions or return default values
|
|
request_interval: Minimum seconds between requests (default: 0.0)
|
|
**kwargs: Additional model-specific parameters
|
|
"""
|
|
self.api_key: str = api_key
|
|
self.base_url: str = base_url
|
|
self.model_name: str = model_name
|
|
self.max_retries: int = max_retries
|
|
self.raise_exception: bool = raise_exception
|
|
self.request_interval: float = request_interval
|
|
self.kwargs: dict = kwargs
|
|
|
|
self._last_request_time: float = 0.0
|
|
self._request_lock: asyncio.Lock = asyncio.Lock()
|
|
|
|
@staticmethod
|
|
def _accumulate_tool_call_chunk(tool_call, ret_tools: list[ToolCall]):
|
|
"""Assemble incremental tool call chunks into complete ToolCall objects."""
|
|
index = tool_call.index
|
|
|
|
while len(ret_tools) <= index:
|
|
ret_tools.append(ToolCall(index=index))
|
|
|
|
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_tool_calls: list[ToolCall], tools: list[ToolCall]) -> list[dict]:
|
|
"""Validate and serialize tool calls."""
|
|
if not ret_tool_calls:
|
|
return []
|
|
|
|
tool_dict: dict[str, ToolCall] = {x.name: x for x in tools} if tools else {}
|
|
validated_tools = []
|
|
|
|
for tool in ret_tool_calls:
|
|
if tool.name not in tool_dict:
|
|
continue
|
|
|
|
if not tool.sanitize_and_check_argument():
|
|
logger.error(f"Invalid JSON arguments in {tool.name}: {tool.arguments}")
|
|
raise ValueError(f"Invalid JSON arguments in {tool.name}: {tool.arguments}")
|
|
|
|
validated_tools.append(tool.simple_output_dump())
|
|
return validated_tools
|
|
|
|
@abstractmethod
|
|
def _build_stream_kwargs(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
log_params: bool = True,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> dict:
|
|
"""Build provider-specific streaming parameters."""
|
|
|
|
async def _stream_chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None,
|
|
stream_kwargs: dict,
|
|
) -> AsyncGenerator[StreamChunk, None]:
|
|
"""Async generator for streaming response chunks."""
|
|
|
|
def _stream_chat_sync(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
stream_kwargs: dict | None = None,
|
|
) -> Generator[StreamChunk, None, None]:
|
|
"""Sync generator for streaming response chunks."""
|
|
|
|
async def stream_chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> AsyncGenerator[StreamChunk, None]:
|
|
"""Stream chat completions with retries and return final message."""
|
|
if self.request_interval > 0:
|
|
async with self._request_lock:
|
|
current_time = time.time()
|
|
elapsed = current_time - self._last_request_time
|
|
if elapsed < self.request_interval:
|
|
await asyncio.sleep(self.request_interval - elapsed)
|
|
self._last_request_time = time.time()
|
|
|
|
async for chunk in self._stream_chat_impl(messages, tools, model_name, **kwargs):
|
|
yield chunk
|
|
|
|
async def _stream_chat_impl(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> AsyncGenerator[StreamChunk, None]:
|
|
"""Stream chat with retry logic."""
|
|
stream_kwargs = self._build_stream_kwargs(messages, tools, model_name=model_name, **kwargs)
|
|
|
|
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
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Stream chat error (model={self.model_name}): {e.args}")
|
|
|
|
if i == self.max_retries - 1:
|
|
if self.raise_exception:
|
|
raise e
|
|
yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e))
|
|
break
|
|
|
|
yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e))
|
|
await asyncio.sleep(i + 1)
|
|
|
|
def stream_chat_sync(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> Generator[StreamChunk, None, None]:
|
|
"""Stream chat completions synchronously with retries."""
|
|
stream_kwargs = self._build_stream_kwargs(messages, tools, model_name=model_name, **kwargs)
|
|
|
|
for i in range(self.max_retries):
|
|
try:
|
|
yield from self._stream_chat_sync(messages=messages, tools=tools, stream_kwargs=stream_kwargs)
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Stream chat sync error (model={self.model_name}): {e.args}")
|
|
|
|
if i == self.max_retries - 1:
|
|
if self.raise_exception:
|
|
raise e
|
|
yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e))
|
|
break
|
|
|
|
yield StreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e))
|
|
time.sleep(i + 1)
|
|
|
|
async def _chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
enable_stream_print: bool = False,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> Message:
|
|
"""Aggregate 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, model_name=model_name, **kwargs)
|
|
async for stream_chunk in self._stream_chat(messages=messages, tools=tools, stream_kwargs=stream_kwargs):
|
|
if stream_chunk.chunk_type is ChunkEnum.USAGE:
|
|
if enable_stream_print:
|
|
print(
|
|
f"\n<usage>{json.dumps(stream_chunk.chunk, ensure_ascii=False, indent=2)}</usage>",
|
|
flush=True,
|
|
)
|
|
|
|
elif stream_chunk.chunk_type is ChunkEnum.THINK:
|
|
if enable_stream_print:
|
|
if not state["enter_think"]:
|
|
state["enter_think"] = True
|
|
print("<think>\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</think>", 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<tool>{json.dumps(stream_chunk.chunk, ensure_ascii=False, indent=2)}</tool>", flush=True)
|
|
state["tool_calls"].append(stream_chunk.chunk)
|
|
|
|
elif stream_chunk.chunk_type is ChunkEnum.ERROR:
|
|
if enable_stream_print:
|
|
print(f"\n<error>{stream_chunk.chunk}</error>", flush=True)
|
|
|
|
return Message(
|
|
role=Role.ASSISTANT,
|
|
reasoning_content=state["reasoning_content"],
|
|
content=state["answer_content"],
|
|
tool_calls=state["tool_calls"],
|
|
)
|
|
|
|
def _chat_sync(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
enable_stream_print: bool = False,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> Message:
|
|
"""Aggregate full response synchronously 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, model_name=model_name, **kwargs)
|
|
for stream_chunk in self._stream_chat_sync(messages=messages, tools=tools, stream_kwargs=stream_kwargs):
|
|
if stream_chunk.chunk_type is ChunkEnum.USAGE:
|
|
if enable_stream_print:
|
|
print(
|
|
f"\n<usage>{json.dumps(stream_chunk.chunk, ensure_ascii=False, indent=2)}</usage>",
|
|
flush=True,
|
|
)
|
|
|
|
elif stream_chunk.chunk_type is ChunkEnum.THINK:
|
|
if enable_stream_print:
|
|
if not state["enter_think"]:
|
|
state["enter_think"] = True
|
|
print("<think>\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</think>", 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<tool>{json.dumps(stream_chunk.chunk, ensure_ascii=False, indent=2)}</tool>", flush=True)
|
|
state["tool_calls"].append(stream_chunk.chunk)
|
|
|
|
elif stream_chunk.chunk_type is ChunkEnum.ERROR:
|
|
if enable_stream_print:
|
|
print(f"\n<error>{stream_chunk.chunk}</error>", flush=True)
|
|
|
|
return Message(
|
|
role=Role.ASSISTANT,
|
|
reasoning_content=state["reasoning_content"],
|
|
content=state["answer_content"],
|
|
tool_calls=state["tool_calls"],
|
|
)
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
enable_stream_print: bool = False,
|
|
callback_fn: Callable[[Message], Any] | None = None,
|
|
default_value: Any = None,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> Message | Any:
|
|
"""Chat completion with retries and error handling."""
|
|
if self.request_interval > 0:
|
|
async with self._request_lock:
|
|
current_time = time.time()
|
|
elapsed = current_time - self._last_request_time
|
|
if elapsed < self.request_interval:
|
|
await asyncio.sleep(self.request_interval - elapsed)
|
|
self._last_request_time = time.time()
|
|
|
|
return await self._chat_impl(
|
|
messages,
|
|
tools,
|
|
enable_stream_print,
|
|
callback_fn,
|
|
default_value,
|
|
model_name,
|
|
**kwargs,
|
|
)
|
|
|
|
async def _chat_impl(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
enable_stream_print: bool = False,
|
|
callback_fn: Callable[[Message], Any] | None = None,
|
|
default_value: Any = None,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> Message | Any:
|
|
"""Chat with retry and error handling logic."""
|
|
effective_model = model_name if model_name is not None else self.model_name
|
|
|
|
for i in range(self.max_retries):
|
|
try:
|
|
result = await self._chat(
|
|
messages=messages,
|
|
tools=tools,
|
|
enable_stream_print=enable_stream_print,
|
|
model_name=model_name,
|
|
**kwargs,
|
|
)
|
|
return callback_fn(result) if callback_fn else result
|
|
|
|
except Exception as e:
|
|
error_message = str(e.args[0]) if e.args else str(e)
|
|
is_inappropriate_content = "inappropriate content" in error_message.lower()
|
|
is_rate_limit_error = (
|
|
"request rate increased too quickly" in error_message.lower()
|
|
or "exceeded your current quota" in error_message.lower()
|
|
or "insufficient_quota" in error_message.lower()
|
|
)
|
|
|
|
if is_inappropriate_content:
|
|
logger.error(f"Inappropriate content detected (model={effective_model})")
|
|
logger.error("=" * 80)
|
|
for idx, msg in enumerate(messages):
|
|
logger.error(f"Message {idx + 1} [role={msg.role}]:")
|
|
logger.error(f"Content: {msg.content}")
|
|
if msg.reasoning_content:
|
|
logger.error(f"Reasoning: {msg.reasoning_content}")
|
|
if msg.tool_calls:
|
|
logger.error(f"Tool calls: {msg.tool_calls}")
|
|
logger.error("-" * 80)
|
|
logger.error("=" * 80)
|
|
return Message(role=Role.ASSISTANT, content="")
|
|
|
|
if is_rate_limit_error:
|
|
logger.warning(
|
|
f"Rate limit hit (model={effective_model}), sleeping 60s (attempt {i + 1}/{self.max_retries})",
|
|
)
|
|
await asyncio.sleep(60)
|
|
continue
|
|
|
|
logger.exception(f"Chat error (model={effective_model}): {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 chat_sync(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[ToolCall] | None = None,
|
|
enable_stream_print: bool = False,
|
|
callback_fn: Callable[[Message], Any] | None = None,
|
|
default_value: Any = None,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> Message | Any:
|
|
"""Chat completion synchronously with retries and error handling."""
|
|
effective_model = model_name if model_name is not None else self.model_name
|
|
|
|
for i in range(self.max_retries):
|
|
try:
|
|
result = self._chat_sync(
|
|
messages=messages,
|
|
tools=tools,
|
|
enable_stream_print=enable_stream_print,
|
|
model_name=model_name,
|
|
**kwargs,
|
|
)
|
|
return callback_fn(result) if callback_fn else result
|
|
|
|
except Exception as e:
|
|
error_message = str(e.args[0]) if e.args else str(e)
|
|
is_inappropriate_content = "inappropriate content" in error_message.lower()
|
|
is_rate_limit_error = (
|
|
"request rate increased too quickly" in error_message.lower()
|
|
or "exceeded your current quota" in error_message.lower()
|
|
or "insufficient_quota" in error_message.lower()
|
|
)
|
|
|
|
if is_inappropriate_content:
|
|
logger.error(f"Inappropriate content detected (model={effective_model})")
|
|
logger.error("=" * 80)
|
|
for idx, msg in enumerate(messages):
|
|
logger.error(f"Message {idx + 1} [role={msg.role}]:")
|
|
logger.error(f"Content: {msg.content}")
|
|
if msg.reasoning_content:
|
|
logger.error(f"Reasoning: {msg.reasoning_content}")
|
|
if msg.tool_calls:
|
|
logger.error(f"Tool calls: {msg.tool_calls}")
|
|
logger.error("-" * 80)
|
|
logger.error("=" * 80)
|
|
return Message(role=Role.ASSISTANT, content="")
|
|
|
|
if is_rate_limit_error:
|
|
logger.warning(
|
|
f"Rate limit hit (model={effective_model}), sleeping 60s (attempt {i + 1}/{self.max_retries})",
|
|
)
|
|
time.sleep(60)
|
|
continue
|
|
|
|
logger.exception(f"Chat sync error (model={effective_model}): {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 simple_request(
|
|
self,
|
|
prompt: str,
|
|
model_name: str | None = None,
|
|
callback_fn: Callable[[Message], Any] | None = None,
|
|
default_value: Any = None,
|
|
**kwargs,
|
|
) -> str:
|
|
"""Make a simple request using the LLM."""
|
|
assistant_message = await self.chat(
|
|
messages=[Message(role=Role.USER, content=prompt)],
|
|
model_name=model_name,
|
|
callback_fn=callback_fn,
|
|
default_value=default_value,
|
|
**kwargs,
|
|
)
|
|
return assistant_message.content
|
|
|
|
async def simple_request_for_json(
|
|
self,
|
|
prompt: str,
|
|
model_name: str | None = None,
|
|
**kwargs,
|
|
) -> dict:
|
|
"""Make a simple request using the LLM and extract JSON."""
|
|
|
|
def extract_fn(message: Message) -> dict:
|
|
return extract_content(message.content)
|
|
|
|
return await self.chat(
|
|
messages=[Message(role=Role.USER, content=prompt)],
|
|
model_name=model_name,
|
|
callback_fn=extract_fn,
|
|
default_value={},
|
|
**kwargs,
|
|
)
|
|
|
|
def start_sync(self):
|
|
"""Synchronously initialize resources."""
|
|
|
|
async def start(self):
|
|
"""Asynchronously initialize resources."""
|
|
|
|
def close_sync(self):
|
|
"""Synchronously release resources and close connections."""
|
|
|
|
async def close(self):
|
|
"""Asynchronously release resources and close connections."""
|