mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
update
This commit is contained in:
parent
ed8dd8d367
commit
ea4170fa09
5 changed files with 366 additions and 44 deletions
|
|
@ -8,36 +8,73 @@ from experiencemaker.schema.app_config import AppConfig
|
|||
|
||||
|
||||
class ConfigParser:
|
||||
"""
|
||||
Configuration parser that handles loading and merging configurations from multiple sources.
|
||||
|
||||
The configuration loading priority (from lowest to highest):
|
||||
1. Default configuration from AppConfig schema
|
||||
2. YAML configuration file
|
||||
3. Command line arguments
|
||||
4. Runtime keyword arguments
|
||||
"""
|
||||
|
||||
def __init__(self, args: list):
|
||||
# step1: default config
|
||||
"""
|
||||
Initialize the configuration parser with command line arguments.
|
||||
|
||||
Args:
|
||||
args: List of command line arguments in dotlist format (e.g., ['key=value'])
|
||||
"""
|
||||
# Step 1: Initialize with default configuration from AppConfig schema
|
||||
self.app_config: DictConfig = OmegaConf.structured(AppConfig)
|
||||
|
||||
# step2: load from config yaml file
|
||||
# Step 2: Load configuration from YAML file
|
||||
# First, parse CLI arguments to check if custom config path is specified
|
||||
cli_config: DictConfig = OmegaConf.from_dotlist(args)
|
||||
temp_config: AppConfig = OmegaConf.to_object(OmegaConf.merge(self.app_config, cli_config))
|
||||
|
||||
# Determine config file path: either from CLI args or use predefined config
|
||||
if temp_config.config_path:
|
||||
# Use custom config path if provided
|
||||
config_path = Path(temp_config.config_path)
|
||||
else:
|
||||
# Use predefined config name from the config directory
|
||||
pre_defined_config = temp_config.pre_defined_config
|
||||
if not pre_defined_config.endswith(".yaml"):
|
||||
pre_defined_config += ".yaml"
|
||||
config_path = Path(__file__).parent / pre_defined_config
|
||||
|
||||
logger.info(f"load config from path={config_path}")
|
||||
yaml_config = OmegaConf.load(config_path)
|
||||
# Merge YAML config with default config
|
||||
self.app_config = OmegaConf.merge(self.app_config, yaml_config)
|
||||
|
||||
# merge cli config
|
||||
# Step 3: Merge CLI arguments (highest priority)
|
||||
self.app_config = OmegaConf.merge(self.app_config, cli_config)
|
||||
|
||||
# Log the final merged configuration
|
||||
app_config_dict = OmegaConf.to_container(self.app_config, resolve=True)
|
||||
logger.info(f"app_config=\n{json.dumps(app_config_dict, indent=2, ensure_ascii=False)}")
|
||||
|
||||
def get_app_config(self, **kwargs) -> AppConfig:
|
||||
"""
|
||||
Get the application configuration with optional runtime overrides.
|
||||
|
||||
Args:
|
||||
**kwargs: Additional configuration parameters to override at runtime
|
||||
|
||||
Returns:
|
||||
AppConfig: The final application configuration object
|
||||
"""
|
||||
# Create a copy of the current configuration
|
||||
app_config = self.app_config.copy()
|
||||
|
||||
# Apply runtime overrides if provided
|
||||
if kwargs:
|
||||
# Convert kwargs to dotlist format for OmegaConf
|
||||
kwargs_list = [f"{k}={v}" for k, v in kwargs.items()]
|
||||
update_config = OmegaConf.from_dotlist(kwargs_list)
|
||||
app_config = OmegaConf.merge(app_config, update_config)
|
||||
|
||||
# Convert OmegaConf DictConfig to structured AppConfig object
|
||||
return OmegaConf.to_object(app_config)
|
||||
|
|
|
|||
|
|
@ -8,42 +8,96 @@ from experiencemaker.schema.vector_node import VectorNode
|
|||
|
||||
|
||||
class BaseEmbeddingModel(BaseModel, ABC):
|
||||
model_name: str = Field(default=..., description="model name")
|
||||
dimensions: int = Field(default=..., description="dimensions")
|
||||
max_retries: int = Field(default=3, description="max retries")
|
||||
raise_exception: bool = Field(default=True, description="raise exception")
|
||||
max_batch_size: int = Field(default=10, description="text-embedding-v4 batch size should not be larger than 10")
|
||||
"""
|
||||
Abstract base class for embedding models.
|
||||
|
||||
This class provides a common interface for various embedding model implementations,
|
||||
including retry logic, error handling, and batch processing capabilities.
|
||||
"""
|
||||
# Model configuration fields
|
||||
model_name: str = Field(default=..., description="Name of the embedding model")
|
||||
dimensions: int = Field(default=..., description="Dimensionality of the embedding vectors")
|
||||
max_retries: int = Field(default=3, description="Maximum number of retry attempts on failure")
|
||||
raise_exception: bool = Field(default=True, description="Whether to raise exceptions after max retries")
|
||||
max_batch_size: int = Field(default=10, description="Maximum batch size for processing (text-embedding-v4 should not exceed 10)")
|
||||
|
||||
def _get_embeddings(self, input_text: str | List[str]):
|
||||
"""
|
||||
Abstract method to get embeddings from the model.
|
||||
|
||||
This method must be implemented by concrete subclasses to provide
|
||||
the actual embedding functionality.
|
||||
|
||||
Args:
|
||||
input_text: Single text string or list of text strings to embed
|
||||
|
||||
Returns:
|
||||
Embedding vector(s) corresponding to the input text(s)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_embeddings(self, input_text: str | List[str]):
|
||||
"""
|
||||
Get embeddings with retry logic and error handling.
|
||||
|
||||
This method wraps the _get_embeddings method with automatic retry
|
||||
functionality in case of failures.
|
||||
|
||||
Args:
|
||||
input_text: Single text string or list of text strings to embed
|
||||
|
||||
Returns:
|
||||
Embedding vector(s) or None if all retries failed and raise_exception is False
|
||||
"""
|
||||
# Retry loop with exponential backoff potential
|
||||
for i in range(self.max_retries):
|
||||
try:
|
||||
return self._get_embeddings(input_text)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}")
|
||||
# If this is the last retry and raise_exception is True, re-raise the exception
|
||||
if i == self.max_retries - 1 and self.raise_exception:
|
||||
raise e
|
||||
|
||||
# Return None if all retries failed and raise_exception is False
|
||||
return None
|
||||
|
||||
def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]):
|
||||
"""
|
||||
Generate embeddings for VectorNode objects and update their vector fields.
|
||||
|
||||
This method handles both single nodes and lists of nodes, with automatic
|
||||
batching for efficient processing of large node lists.
|
||||
|
||||
Args:
|
||||
nodes: Single VectorNode or list of VectorNode objects to embed
|
||||
|
||||
Returns:
|
||||
The same node(s) with updated vector fields containing embeddings
|
||||
|
||||
Raises:
|
||||
RuntimeError: If unsupported node type is provided
|
||||
"""
|
||||
# Handle single VectorNode
|
||||
if isinstance(nodes, VectorNode):
|
||||
nodes.vector = self.get_embeddings(nodes.content)
|
||||
return nodes
|
||||
|
||||
# Handle list of VectorNodes with batch processing
|
||||
elif isinstance(nodes, list):
|
||||
|
||||
# Process nodes in batches to respect max_batch_size limits
|
||||
embeddings = [emb for i in range(0, len(nodes), self.max_batch_size) for emb in
|
||||
self.get_embeddings(input_text=[node.content for node in nodes[i:i + self.max_batch_size]])]
|
||||
|
||||
# Validate that we got the expected number of embeddings
|
||||
if len(embeddings) != len(nodes):
|
||||
logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}")
|
||||
else:
|
||||
# Assign embeddings to corresponding nodes
|
||||
for node, embedding in zip(nodes, embeddings):
|
||||
node.vector = embedding
|
||||
return nodes
|
||||
|
||||
else:
|
||||
else:
|
||||
raise RuntimeError(f"unsupported type={type(nodes)}")
|
||||
|
|
|
|||
|
|
@ -11,19 +11,53 @@ from experiencemaker.embedding_model.base_embedding_model import BaseEmbeddingMo
|
|||
|
||||
@EMBEDDING_MODEL_REGISTRY.register("openai_compatible")
|
||||
class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
|
||||
api_key: str = Field(default_factory=lambda: os.getenv("EMBEDDING_API_KEY"), description="api key")
|
||||
base_url: str = Field(default_factory=lambda: os.getenv("EMBEDDING_BASE_URL"), description="base url")
|
||||
model_name: str = Field(default="", description="model name")
|
||||
dimensions: int = Field(default=1024, description="dimensions")
|
||||
encoding_format: Literal["float", "base64"] = Field(default="float", description="encoding_format")
|
||||
"""
|
||||
OpenAI-compatible embedding model implementation.
|
||||
|
||||
This class provides an implementation of BaseEmbeddingModel that works with
|
||||
OpenAI-compatible embedding APIs, including OpenAI's official API and
|
||||
other services that follow the same interface.
|
||||
"""
|
||||
# API configuration fields
|
||||
api_key: str = Field(default_factory=lambda: os.getenv("EMBEDDING_API_KEY"), description="API key for authentication")
|
||||
base_url: str = Field(default_factory=lambda: os.getenv("EMBEDDING_BASE_URL"), description="Base URL for the API endpoint")
|
||||
model_name: str = Field(default="", description="Name of the embedding model to use")
|
||||
dimensions: int = Field(default=1024, description="Dimensionality of the embedding vectors")
|
||||
encoding_format: Literal["float", "base64"] = Field(default="float", description="Encoding format for embeddings")
|
||||
|
||||
# Private OpenAI client instance
|
||||
_client: OpenAI = PrivateAttr()
|
||||
|
||||
@model_validator(mode="after")
|
||||
def init_client(self):
|
||||
"""
|
||||
Initialize the OpenAI client after model validation.
|
||||
|
||||
This method is called automatically after Pydantic model validation
|
||||
to set up the OpenAI client with the provided API key and base URL.
|
||||
|
||||
Returns:
|
||||
self: The model instance for method chaining
|
||||
"""
|
||||
self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
|
||||
return self
|
||||
|
||||
def _get_embeddings(self, input_text: str | List[str]):
|
||||
"""
|
||||
Get embeddings from the OpenAI-compatible API.
|
||||
|
||||
This method implements the abstract _get_embeddings method from BaseEmbeddingModel
|
||||
by calling the OpenAI-compatible embeddings API.
|
||||
|
||||
Args:
|
||||
input_text: Single text string or list of text strings to embed
|
||||
|
||||
Returns:
|
||||
Embedding vector(s) corresponding to the input text(s)
|
||||
|
||||
Raises:
|
||||
RuntimeError: If unsupported input type is provided
|
||||
"""
|
||||
completion = self._client.embeddings.create(
|
||||
model=self.model_name,
|
||||
input=input_text,
|
||||
|
|
|
|||
|
|
@ -10,36 +10,113 @@ from experiencemaker.tool.base_tool import BaseTool
|
|||
|
||||
|
||||
class BaseLLM(BaseModel, ABC):
|
||||
model_name: str = Field(...)
|
||||
"""
|
||||
Abstract base class for Large Language Model (LLM) implementations.
|
||||
|
||||
This class defines the common interface and configuration parameters
|
||||
that all LLM implementations should support. It provides a standardized
|
||||
way to interact with different LLM providers while handling common
|
||||
concerns like retries, error handling, and streaming.
|
||||
"""
|
||||
# Core model configuration
|
||||
model_name: str = Field(..., description="Name of the LLM model to use")
|
||||
|
||||
seed: int = Field(default=42)
|
||||
top_p: float | None = Field(default=None)
|
||||
# stream: bool = Field(default=True)
|
||||
stream_options: dict = Field(default={"include_usage": True})
|
||||
temperature: float = Field(default=0.0000001)
|
||||
presence_penalty: float | None = Field(default=None)
|
||||
enable_thinking: bool = Field(default=True, description="whether the current mode is the reasoning model, "
|
||||
"or whether Qwen3's reasoning mode is currently enabled.")
|
||||
tool_choice: Literal["none", "auto", "required"] = Field(default="auto", description="tool choice")
|
||||
parallel_tool_calls: bool = Field(default=True)
|
||||
# Generation parameters
|
||||
seed: int = Field(default=42, description="Random seed for reproducible outputs")
|
||||
top_p: float | None = Field(default=None, description="Top-p (nucleus) sampling parameter")
|
||||
# stream: bool = Field(default=True) # Commented out - streaming is handled per request
|
||||
stream_options: dict = Field(default={"include_usage": True}, description="Options for streaming responses")
|
||||
temperature: float = Field(default=0.0000001, description="Sampling temperature (low for deterministic outputs)")
|
||||
presence_penalty: float | None = Field(default=None, description="Presence penalty to reduce repetition")
|
||||
|
||||
# Model-specific features
|
||||
enable_thinking: bool = Field(default=True, description="Enable reasoning/thinking mode for supported models")
|
||||
|
||||
# Tool usage configuration
|
||||
tool_choice: Literal["none", "auto", "required"] = Field(default="auto", description="Strategy for tool selection")
|
||||
parallel_tool_calls: bool = Field(default=True, description="Allow multiple tool calls in parallel")
|
||||
|
||||
max_retries: int = Field(default=5, description="max retries")
|
||||
raise_exception: bool = Field(default=False, description="raise exception")
|
||||
# Error handling and reliability
|
||||
max_retries: int = Field(default=5, description="Maximum number of retry attempts on failure")
|
||||
raise_exception: bool = Field(default=False, description="Whether to raise exceptions or return default values")
|
||||
|
||||
def stream_chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs):
|
||||
"""
|
||||
Stream chat completions from the LLM.
|
||||
|
||||
This method should yield chunks of the response as they become available,
|
||||
allowing for real-time display of the model's output.
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools the model can use
|
||||
**kwargs: Additional model-specific parameters
|
||||
|
||||
Yields:
|
||||
Chunks of the streaming response with their types
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def stream_print(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs):
|
||||
"""
|
||||
Stream chat completions and print them to console in real-time.
|
||||
|
||||
This is a convenience method for debugging and interactive use,
|
||||
combining streaming with formatted console output.
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools the model can use
|
||||
**kwargs: Additional model-specific parameters
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs) -> Message:
|
||||
"""
|
||||
Internal method to perform a single chat completion.
|
||||
|
||||
This method should be implemented by subclasses to handle the actual
|
||||
communication with the LLM provider. It's called by the public chat()
|
||||
method which adds retry logic and error handling.
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools the model can use
|
||||
**kwargs: Additional model-specific parameters
|
||||
|
||||
Returns:
|
||||
The complete response message from the LLM
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def chat(self, messages: List[Message], tools: List[BaseTool] = None, callback_fn: Callable = None,
|
||||
default_value=None, **kwargs):
|
||||
"""
|
||||
Perform a chat completion with retry logic and error handling.
|
||||
|
||||
This is the main public interface for chat completions. It wraps the
|
||||
internal _chat() method with robust error handling, exponential backoff,
|
||||
and optional callback processing.
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools the model can use
|
||||
callback_fn: Optional callback to process the response message
|
||||
default_value: Value to return if all retries fail (when raise_exception=False)
|
||||
**kwargs: Additional model-specific parameters
|
||||
|
||||
Returns:
|
||||
The response message (possibly processed by callback_fn) or default_value
|
||||
|
||||
Raises:
|
||||
Exception: If raise_exception=True and all retries fail
|
||||
"""
|
||||
for i in range(self.max_retries):
|
||||
try:
|
||||
# Attempt to get response from the model
|
||||
message: Message = self._chat(messages, tools, **kwargs)
|
||||
|
||||
# Apply callback function if provided
|
||||
if callback_fn:
|
||||
return callback_fn(message)
|
||||
else:
|
||||
|
|
@ -47,8 +124,11 @@ class BaseLLM(BaseModel, ABC):
|
|||
|
||||
except Exception as e:
|
||||
logger.exception(f"chat with model={self.model_name} encounter error with e={e.args}")
|
||||
|
||||
# Exponential backoff: wait longer after each failure
|
||||
time.sleep(1 + i)
|
||||
|
||||
# Handle final retry failure
|
||||
if i == self.max_retries - 1:
|
||||
if self.raise_exception:
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -17,18 +17,58 @@ from experiencemaker.tool.base_tool import BaseTool
|
|||
|
||||
@LLM_REGISTRY.register("openai_compatible")
|
||||
class OpenAICompatibleBaseLLM(BaseLLM):
|
||||
api_key: str = Field(default_factory=lambda: os.getenv("LLM_API_KEY"), description="api key")
|
||||
base_url: str = Field(default_factory=lambda: os.getenv("LLM_BASE_URL"), description="base url")
|
||||
_client: OpenAI = PrivateAttr()
|
||||
"""
|
||||
OpenAI-compatible LLM implementation supporting streaming and tool calls.
|
||||
|
||||
This class implements the BaseLLM interface for OpenAI-compatible APIs,
|
||||
including support for:
|
||||
- Streaming responses with different chunk types (thinking, answer, tools)
|
||||
- Tool calling with parallel execution
|
||||
- Reasoning/thinking content from supported models
|
||||
- Robust error handling and retries
|
||||
"""
|
||||
|
||||
# API configuration
|
||||
api_key: str = Field(default_factory=lambda: os.getenv("LLM_API_KEY"), description="API key for authentication")
|
||||
base_url: str = Field(default_factory=lambda: os.getenv("LLM_BASE_URL"), description="Base URL for the API endpoint")
|
||||
_client: OpenAI = PrivateAttr(description="OpenAI client instance (private)")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def init_client(self):
|
||||
"""
|
||||
Initialize the OpenAI client after model validation.
|
||||
|
||||
This validator runs after all field validation is complete,
|
||||
ensuring we have valid API credentials before creating the client.
|
||||
|
||||
Returns:
|
||||
Self for method chaining
|
||||
"""
|
||||
self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
|
||||
return self
|
||||
|
||||
def stream_chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs):
|
||||
"""
|
||||
Stream chat completions from OpenAI-compatible API.
|
||||
|
||||
This method handles streaming responses and categorizes chunks into different types:
|
||||
- THINK: Reasoning/thinking content from the model
|
||||
- ANSWER: Regular response content
|
||||
- TOOL: Tool calls that need to be executed
|
||||
- USAGE: Token usage statistics
|
||||
- ERROR: Error information
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools available to the model
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Yields:
|
||||
Tuple of (chunk_content, ChunkEnum) for each streaming piece
|
||||
"""
|
||||
for i in range(self.max_retries):
|
||||
try:
|
||||
# Create streaming completion request
|
||||
completion = self._client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=[x.simple_dump() for x in messages],
|
||||
|
|
@ -37,37 +77,47 @@ class OpenAICompatibleBaseLLM(BaseLLM):
|
|||
stream=True,
|
||||
stream_options=self.stream_options,
|
||||
temperature=self.temperature,
|
||||
extra_body={"enable_thinking": self.enable_thinking},
|
||||
extra_body={"enable_thinking": self.enable_thinking}, # Enable reasoning mode
|
||||
tools=[x.simple_dump() for x in tools] if tools else None,
|
||||
tool_choice=self.tool_choice,
|
||||
parallel_tool_calls=self.parallel_tool_calls)
|
||||
|
||||
ret_tools = []
|
||||
is_answering = False
|
||||
# Initialize tool call tracking
|
||||
ret_tools = [] # Accumulate tool calls across chunks
|
||||
is_answering = False # Track when model starts answering
|
||||
|
||||
# Process each chunk in the streaming response
|
||||
for chunk in completion:
|
||||
# Handle chunks without choices (usually usage info)
|
||||
if not chunk.choices:
|
||||
yield chunk.usage, ChunkEnum.USAGE
|
||||
|
||||
else:
|
||||
delta = chunk.choices[0].delta
|
||||
|
||||
# Handle reasoning/thinking content (model's internal thoughts)
|
||||
if hasattr(delta, 'reasoning_content') and delta.reasoning_content is not None:
|
||||
yield delta.reasoning_content, ChunkEnum.THINK
|
||||
|
||||
else:
|
||||
# Mark transition from thinking to answering
|
||||
if not is_answering:
|
||||
is_answering = True
|
||||
|
||||
# Handle regular response content
|
||||
if delta.content is not None:
|
||||
yield delta.content, ChunkEnum.ANSWER
|
||||
|
||||
# Handle tool calls (function calling)
|
||||
if delta.tool_calls is not None:
|
||||
for tool_call in delta.tool_calls:
|
||||
index = tool_call.index
|
||||
|
||||
# Ensure we have enough tool call slots
|
||||
while len(ret_tools) <= index:
|
||||
ret_tools.append(ToolCall(index=index))
|
||||
|
||||
# Accumulate tool call information across chunks
|
||||
if tool_call.id:
|
||||
ret_tools[index].id += tool_call.id
|
||||
|
||||
|
|
@ -77,28 +127,49 @@ class OpenAICompatibleBaseLLM(BaseLLM):
|
|||
if tool_call.function and tool_call.function.arguments:
|
||||
ret_tools[index].arguments += tool_call.function.arguments
|
||||
|
||||
# Yield completed tool calls after streaming finishes
|
||||
if ret_tools:
|
||||
tool_dict = {x.name: x for x in tools}
|
||||
tool_dict = {x.name: x for x in tools} if tools else {}
|
||||
for tool in ret_tools:
|
||||
# Only yield tool calls that correspond to available tools
|
||||
if tool.name not in tool_dict:
|
||||
continue
|
||||
|
||||
yield tool, ChunkEnum.TOOL
|
||||
|
||||
return
|
||||
return # Success - exit retry loop
|
||||
|
||||
except Exception as e:
|
||||
except Exception as e:
|
||||
logger.exception(f"stream chat with model={self.model_name} encounter error with e={e.args}")
|
||||
|
||||
# Handle retry logic
|
||||
if i == self.max_retries - 1 and self.raise_exception:
|
||||
raise e
|
||||
else:
|
||||
yield e.args, ChunkEnum.ERROR
|
||||
|
||||
def _chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs) -> Message:
|
||||
reasoning_content = ""
|
||||
answer_content = ""
|
||||
tool_calls = []
|
||||
"""
|
||||
Perform a complete chat completion by aggregating streaming chunks.
|
||||
|
||||
This method consumes the entire streaming response and combines all
|
||||
chunks into a single Message object. It separates reasoning content,
|
||||
regular answer content, and tool calls.
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools available to the model
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Complete Message with all content aggregated
|
||||
"""
|
||||
# Initialize content accumulators
|
||||
reasoning_content = "" # Model's internal reasoning
|
||||
answer_content = "" # Final response content
|
||||
tool_calls = [] # List of tool calls to execute
|
||||
|
||||
# Consume streaming response and aggregate chunks by type
|
||||
for chunk, chunk_enum in self.stream_chat(messages, tools, **kwargs):
|
||||
if chunk_enum is ChunkEnum.THINK:
|
||||
reasoning_content += chunk
|
||||
|
|
@ -108,58 +179,104 @@ class OpenAICompatibleBaseLLM(BaseLLM):
|
|||
|
||||
elif chunk_enum is ChunkEnum.TOOL:
|
||||
tool_calls.append(chunk)
|
||||
|
||||
# Note: USAGE and ERROR chunks are ignored in non-streaming mode
|
||||
|
||||
# Construct complete response message
|
||||
return Message(role=Role.ASSISTANT,
|
||||
reasoning_content=reasoning_content,
|
||||
content=answer_content,
|
||||
tool_calls=tool_calls)
|
||||
|
||||
def stream_print(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs):
|
||||
enter_think = False
|
||||
enter_answer = False
|
||||
"""
|
||||
Stream chat completions with formatted console output.
|
||||
|
||||
This method provides a real-time view of the model's response,
|
||||
with different formatting for different types of content:
|
||||
- Thinking content is wrapped in <think></think> tags
|
||||
- Answer content is printed directly
|
||||
- Tool calls are formatted as JSON
|
||||
- Usage statistics and errors are clearly marked
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tools available to the model
|
||||
**kwargs: Additional parameters
|
||||
"""
|
||||
# Track which sections we've entered for proper formatting
|
||||
enter_think = False # Whether we've started printing thinking content
|
||||
enter_answer = False # Whether we've started printing answer content
|
||||
|
||||
# Process each streaming chunk with appropriate formatting
|
||||
for chunk, chunk_enum in self.stream_chat(messages, tools, **kwargs):
|
||||
if chunk_enum is ChunkEnum.USAGE:
|
||||
# Display token usage statistics
|
||||
if isinstance(chunk, CompletionUsage):
|
||||
print(f"\n<usage>{chunk.model_dump_json(indent=2)}</usage>")
|
||||
else:
|
||||
print(f"\n<usage>{chunk}</usage>")
|
||||
|
||||
elif chunk_enum is ChunkEnum.THINK:
|
||||
# Format thinking/reasoning content
|
||||
if not enter_think:
|
||||
enter_think = True
|
||||
print("<think>\n", end="")
|
||||
print(chunk, end="")
|
||||
|
||||
elif chunk_enum is ChunkEnum.ANSWER:
|
||||
# Format regular answer content
|
||||
if not enter_answer:
|
||||
enter_answer = True
|
||||
# Close thinking section if we were in it
|
||||
if enter_think:
|
||||
print("\n</think>")
|
||||
print(chunk, end="")
|
||||
|
||||
elif chunk_enum is ChunkEnum.TOOL:
|
||||
# Format tool calls as structured JSON
|
||||
assert isinstance(chunk, ToolCall)
|
||||
print(f"\n<tool>{chunk.model_dump_json(indent=2)}</tool>", end="")
|
||||
|
||||
elif chunk_enum is ChunkEnum.ERROR:
|
||||
# Display error information
|
||||
print(f"\n<error>{chunk}</error>", end="")
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Demo function to test the OpenAI-compatible LLM implementation.
|
||||
|
||||
This function demonstrates:
|
||||
1. Basic chat without tools
|
||||
2. Chat with tool usage (search and code tools)
|
||||
3. Real-time streaming output formatting
|
||||
"""
|
||||
from experiencemaker.tool.dashscope_search_tool import DashscopeSearchTool
|
||||
from experiencemaker.tool.code_tool import CodeTool
|
||||
from experiencemaker.enumeration.role import Role
|
||||
|
||||
# Load environment variables for API credentials
|
||||
load_dotenv()
|
||||
|
||||
# Initialize the LLM with a specific model
|
||||
model_name = "qwen-max-2025-01-25"
|
||||
llm = OpenAICompatibleBaseLLM(model_name=model_name)
|
||||
|
||||
# Set up available tools
|
||||
tools: List[BaseTool] = [DashscopeSearchTool(), CodeTool()]
|
||||
|
||||
# Test 1: Simple greeting without tools
|
||||
print("=== Test 1: Simple Chat ===")
|
||||
llm.stream_print([Message(role=Role.USER, content="hello")], [])
|
||||
print("=" * 20)
|
||||
|
||||
print("\n" + "=" * 20)
|
||||
|
||||
# Test 2: Complex query that might use tools
|
||||
print("\n=== Test 2: Chat with Tools ===")
|
||||
llm.stream_print([Message(role=Role.USER, content="What's the weather like in Beijing today?")], tools)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# launch with: python -m experiencemaker.model.openai_compatible_llm
|
||||
# Launch with: python -m experiencemaker.llm.openai_compatible_llm
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue