first commit
This commit is contained in:
@@ -0,0 +1,324 @@
|
||||
"""
|
||||
Agent registration mechanism.
|
||||
Used to manage multiple Agent instances and supports selecting different Agents by model name.
|
||||
"""
|
||||
|
||||
from typing import Dict, Optional, Type, List, Any
|
||||
from .BaseAgent import BaseAgent
|
||||
from config.config import get_config
|
||||
import threading
|
||||
|
||||
|
||||
class AgentRegistry:
|
||||
"""Agent registry, managing multiple Agent instances."""
|
||||
|
||||
def __init__(self):
|
||||
self._agents: Dict[str, BaseAgent] = {}
|
||||
self._agent_classes: Dict[str, Type[BaseAgent]] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def register_agent_class(self, model_name: str, agent_class: Type[BaseAgent]):
|
||||
"""
|
||||
Register an Agent class.
|
||||
|
||||
Args:
|
||||
model_name: Model name, used for the model parameter in the API
|
||||
agent_class: Agent class, inheriting from BaseAgent
|
||||
"""
|
||||
with self._lock:
|
||||
self._agent_classes[model_name] = agent_class
|
||||
|
||||
def _create_agent_instance(
|
||||
self,
|
||||
model_name: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
context_file: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> BaseAgent:
|
||||
"""
|
||||
Internal method: create an Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
name: Agent name
|
||||
description: Agent description
|
||||
context_file: Context file path
|
||||
**kwargs: Other parameters
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
if model_name not in self._agent_classes:
|
||||
raise ValueError(f"Unknown model: {model_name}")
|
||||
|
||||
agent_class = self._agent_classes[model_name]
|
||||
|
||||
# Get configuration.
|
||||
config = get_config()
|
||||
|
||||
# Use default values or passed-in parameters.
|
||||
agent_name = name or f"{model_name}-agent"
|
||||
agent_description = description or f"Agent instance for {model_name}"
|
||||
|
||||
# Create the Agent instance.
|
||||
agent = agent_class(
|
||||
name=agent_name,
|
||||
description=agent_description,
|
||||
llm_interface=config.BASIC_INTERFACE,
|
||||
context_file=context_file,
|
||||
model_name=model_name, # Pass model_name to the Agent.
|
||||
**kwargs
|
||||
)
|
||||
|
||||
return agent
|
||||
|
||||
def get_or_create_agent(
|
||||
self,
|
||||
model_name: str,
|
||||
**kwargs
|
||||
) -> BaseAgent:
|
||||
"""
|
||||
Get or create an Agent instance (singleton pattern).
|
||||
Ensure each model_name corresponds to only one Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
**kwargs: Creation parameters
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
with self._lock:
|
||||
if model_name not in self._agents:
|
||||
self._agents[model_name] = self._create_agent_instance(model_name, **kwargs)
|
||||
return self._agents[model_name]
|
||||
|
||||
def create_agent(
|
||||
self,
|
||||
model_name: str,
|
||||
force_new: bool = False,
|
||||
**kwargs
|
||||
) -> BaseAgent:
|
||||
"""
|
||||
Create an Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
force_new: Whether to force creation of a new instance, replacing the existing instance
|
||||
**kwargs: Other parameters
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
with self._lock:
|
||||
if force_new or model_name not in self._agents:
|
||||
self._agents[model_name] = self._create_agent_instance(model_name, **kwargs)
|
||||
return self._agents[model_name]
|
||||
|
||||
def get_agent(self, model_name: str) -> Optional[BaseAgent]:
|
||||
"""
|
||||
Get an already-created Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
|
||||
Returns:
|
||||
Agent instance or None
|
||||
"""
|
||||
return self._agents.get(model_name)
|
||||
|
||||
def list_models(self) -> List[str]:
|
||||
"""
|
||||
List all registered model names.
|
||||
|
||||
Returns:
|
||||
List of model names
|
||||
"""
|
||||
return list(self._agent_classes.keys())
|
||||
|
||||
def list_agents(self) -> List[str]:
|
||||
"""
|
||||
List all created Agent instances.
|
||||
|
||||
Returns:
|
||||
List of model names for Agent instances
|
||||
"""
|
||||
return list(self._agents.keys())
|
||||
|
||||
def clear_agents(self):
|
||||
"""Clear all Agent instances."""
|
||||
with self._lock:
|
||||
self._agents.clear()
|
||||
|
||||
def remove_agent(self, model_name: str) -> bool:
|
||||
"""
|
||||
Remove an Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
|
||||
Returns:
|
||||
Whether removal succeeded
|
||||
"""
|
||||
with self._lock:
|
||||
if model_name in self._agents:
|
||||
del self._agents[model_name]
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_agent_info(self, model_name: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get Agent information.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
|
||||
Returns:
|
||||
Agent information dictionary or None
|
||||
"""
|
||||
agent = self.get_agent(model_name)
|
||||
if agent:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"name": agent.name,
|
||||
"description": agent.description,
|
||||
"agent_class": agent.__class__.__name__,
|
||||
"toolkit_size": len(agent.toolkit),
|
||||
"session_info": agent.get_session_info(),
|
||||
"is_singleton": True # Mark this as a singleton instance.
|
||||
}
|
||||
return None
|
||||
|
||||
def get_all_agents_info(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
Get information for all Agents.
|
||||
|
||||
Returns:
|
||||
Dictionary of all Agent information
|
||||
"""
|
||||
result = {}
|
||||
with self._lock:
|
||||
for model_name in self._agents:
|
||||
info = self.get_agent_info(model_name)
|
||||
if info:
|
||||
result[model_name] = info
|
||||
return result
|
||||
|
||||
def is_agent_active(self, model_name: str) -> bool:
|
||||
"""
|
||||
Check whether the Agent with the specified model_name has been created.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
|
||||
Returns:
|
||||
Whether it has been created
|
||||
"""
|
||||
return model_name in self._agents
|
||||
|
||||
def get_agent_stats(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get registry statistics.
|
||||
|
||||
Returns:
|
||||
Statistics dictionary
|
||||
"""
|
||||
with self._lock:
|
||||
return {
|
||||
"registered_models": len(self._agent_classes),
|
||||
"active_agents": len(self._agents),
|
||||
"registered_model_list": list(self._agent_classes.keys()),
|
||||
"active_agent_list": list(self._agents.keys())
|
||||
}
|
||||
|
||||
|
||||
# Global Agent registry instance.
|
||||
_global_registry = AgentRegistry()
|
||||
|
||||
|
||||
def get_agent_registry() -> AgentRegistry:
|
||||
"""Get the global Agent registry."""
|
||||
return _global_registry
|
||||
|
||||
|
||||
def register_agent(model_name: str, agent_class: Type[BaseAgent]):
|
||||
"""
|
||||
Convenience function for registering an Agent class.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
agent_class: Agent class
|
||||
"""
|
||||
_global_registry.register_agent_class(model_name, agent_class)
|
||||
|
||||
|
||||
def get_agent(model_name: str, **kwargs) -> BaseAgent:
|
||||
"""
|
||||
Convenience function for getting or creating an Agent instance (singleton pattern).
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
**kwargs: Creation parameters, used only during first creation
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
return _global_registry.get_or_create_agent(model_name, **kwargs)
|
||||
|
||||
|
||||
def get_existing_agent(model_name: str) -> Optional[BaseAgent]:
|
||||
"""
|
||||
Convenience function for getting an existing Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
|
||||
Returns:
|
||||
Agent instance or None
|
||||
"""
|
||||
return _global_registry.get_agent(model_name)
|
||||
|
||||
|
||||
def create_new_agent(model_name: str, **kwargs) -> BaseAgent:
|
||||
"""
|
||||
Convenience function for forcing creation of a new Agent instance.
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
**kwargs: Creation parameters
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
return _global_registry.create_agent(model_name, force_new=True, **kwargs)
|
||||
|
||||
|
||||
def list_available_models() -> List[str]:
|
||||
"""
|
||||
Convenience function for listing all available models.
|
||||
|
||||
Returns:
|
||||
List of model names
|
||||
"""
|
||||
return _global_registry.list_models()
|
||||
|
||||
|
||||
def get_registry_stats() -> Dict[str, Any]:
|
||||
"""
|
||||
Convenience function for getting registry statistics.
|
||||
|
||||
Returns:
|
||||
Statistics dictionary
|
||||
"""
|
||||
return _global_registry.get_agent_stats()
|
||||
|
||||
|
||||
def clear_all_agents():
|
||||
"""
|
||||
Convenience function for clearing all Agent instances.
|
||||
"""
|
||||
_global_registry.clear_agents()
|
||||
# Also clear BaseAgent's class-level instance cache.
|
||||
from .BaseAgent import BaseAgent
|
||||
BaseAgent.clear_instances()
|
||||
@@ -0,0 +1,473 @@
|
||||
"""
|
||||
BaseAgent is the base class for all agents. It defines the basic agent interface and common functionality.
|
||||
All concrete agent implementations should inherit from this class and implement the abstract methods.
|
||||
|
||||
BaseAgent provides the following features:
|
||||
- Singleton pattern
|
||||
- Conversation history management
|
||||
- SketchPad management
|
||||
- Toolkit management
|
||||
"""
|
||||
|
||||
from bootstrap_env import load_project_env
|
||||
|
||||
load_project_env()
|
||||
|
||||
from typing import (
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Generator,
|
||||
Sequence,
|
||||
Tuple,
|
||||
AsyncGenerator,
|
||||
Any,
|
||||
)
|
||||
from abc import ABC, abstractmethod
|
||||
from SimpleLLMFunc import llm_chat, OpenAICompatible # type: ignore
|
||||
import threading
|
||||
from context.conversation_manager import get_current_context, get_current_sketch_pad
|
||||
from context.schemas import Message
|
||||
from react_stream import extract_output_text, is_response_yield
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
|
||||
|
||||
class BaseAgent(ABC):
|
||||
"""
|
||||
Agent base class, defining the basic agent interface and common functionality.
|
||||
|
||||
All concrete agent implementations should inherit from this class and implement the abstract methods.
|
||||
"""
|
||||
|
||||
# Class-level instance cache, ensuring a singleton for each Agent subclass.
|
||||
_class_instances: Dict[str, "BaseAgent"] = {}
|
||||
_class_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def get_instance(
|
||||
cls,
|
||||
model_name: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
llm_interface: Optional[OpenAICompatible] = None,
|
||||
**kwargs,
|
||||
) -> "BaseAgent":
|
||||
"""
|
||||
Class method for obtaining an Agent instance (singleton pattern).
|
||||
|
||||
Args:
|
||||
model_name: Model name
|
||||
name: Agent name
|
||||
description: Agent description
|
||||
llm_interface: LLM interface
|
||||
**kwargs: Other parameters
|
||||
|
||||
Returns:
|
||||
Agent instance
|
||||
"""
|
||||
with cls._class_lock:
|
||||
# Use the class name and model_name as the unique identifier.
|
||||
instance_key = f"{cls.__name__}:{model_name}"
|
||||
|
||||
if instance_key not in cls._class_instances:
|
||||
if not llm_interface:
|
||||
# If llm_interface is not provided, try to obtain it from the configuration.
|
||||
from config.config import get_config
|
||||
|
||||
config = get_config()
|
||||
llm_interface = config.BASIC_INTERFACE
|
||||
|
||||
instance_name = name or f"{model_name}-agent"
|
||||
instance_description = description or f"Agent instance for {model_name}"
|
||||
|
||||
cls._class_instances[instance_key] = cls(
|
||||
name=instance_name,
|
||||
description=instance_description,
|
||||
llm_interface=llm_interface,
|
||||
model_name=model_name,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return cls._class_instances[instance_key]
|
||||
|
||||
@classmethod
|
||||
def clear_instances(cls):
|
||||
"""Clear all instance caches."""
|
||||
with cls._class_lock:
|
||||
cls._class_instances.clear()
|
||||
|
||||
@classmethod
|
||||
def get_all_instances(cls) -> Dict[str, "BaseAgent"]:
|
||||
"""Get all instances."""
|
||||
return cls._class_instances.copy()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
llm_interface: Optional[OpenAICompatible] = None,
|
||||
model_name: Optional[str] = None, # Add the model_name parameter.
|
||||
**kwargs, # Extra parameters that subclasses can handle.
|
||||
):
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.model_name = model_name # Store model_name.
|
||||
self.llm_interface = llm_interface
|
||||
|
||||
if not self.llm_interface:
|
||||
raise ValueError("llm_interface must be provided")
|
||||
|
||||
# Subclasses need to define their own toolkit.
|
||||
self.toolkit = self.get_toolkit()
|
||||
|
||||
# Initialize the chat function.
|
||||
self.chat = llm_chat(
|
||||
llm_interface=self.llm_interface,
|
||||
toolkit=self.toolkit, # type: ignore
|
||||
stream=True,
|
||||
return_mode="raw",
|
||||
enable_event=True,
|
||||
max_tool_calls=2000,
|
||||
timeout=600,
|
||||
temperature=1.0,
|
||||
)(self.chat_impl)
|
||||
|
||||
@abstractmethod
|
||||
def get_toolkit(self) -> Sequence[Any]:
|
||||
"""
|
||||
Get the agent-specific toolkit (abstract method).
|
||||
|
||||
Subclasses must implement this method to define their own toolkit.
|
||||
|
||||
Returns:
|
||||
List of tool functions
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def chat_impl(
|
||||
self,
|
||||
history: List[Dict[str, Any]],
|
||||
query: Any,
|
||||
sketch_pad_summary: str,
|
||||
) -> Generator[Tuple[str, List[Dict[str, Any]]], None, None]:
|
||||
"""
|
||||
Agent conversation implementation logic (abstract method).
|
||||
|
||||
Subclasses must implement this method to define the concrete conversation behavior.
|
||||
|
||||
Args:
|
||||
history: Conversation history
|
||||
query: User query
|
||||
sketch_pad_summary: SketchPad summary
|
||||
|
||||
Returns:
|
||||
Generator yielding (response_chunk, updated_history)
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def run(
|
||||
self, query: Any, raw_user_content: Any = None
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""
|
||||
Run the agent to process the user query (abstract method).
|
||||
|
||||
Args:
|
||||
query: User query
|
||||
raw_user_content: Raw user message content (optional), used for persisting multimodal messages
|
||||
|
||||
Returns:
|
||||
AsyncGenerator yielding response chunks
|
||||
"""
|
||||
pass
|
||||
|
||||
# Common helper methods.
|
||||
def get_sketch_pad_summary(self) -> str:
|
||||
"""Get SketchPad summary information, including all keys and truncated values."""
|
||||
try:
|
||||
sketch_pad = get_current_sketch_pad()
|
||||
if sketch_pad is None:
|
||||
return "SketchPad unavailable: no active conversation context"
|
||||
|
||||
# Get detailed information for all items (including values).
|
||||
all_items = sketch_pad.list_items(include_value=True)
|
||||
|
||||
if not all_items:
|
||||
return "SketchPad is empty: no stored content"
|
||||
|
||||
summary_lines = [f"Current SketchPad state ({len(all_items)} items total):"]
|
||||
|
||||
for item in all_items[:20]: # Limit display to the first 20 items.
|
||||
key = item.key
|
||||
tags = ", ".join(item.tags) if item.tags else "no tags"
|
||||
timestamp = item.timestamp
|
||||
content_type = item.content_type
|
||||
|
||||
# Use the value included in the list item for preview.
|
||||
value_obj = item.value
|
||||
value_str = str(value_obj) if value_obj is not None else ""
|
||||
if len(value_str) > 100:
|
||||
value_preview = value_str[:100] + "..."
|
||||
else:
|
||||
value_preview = value_str
|
||||
|
||||
value_preview = value_preview.replace("\n", "\\n")
|
||||
|
||||
summary_lines.append(
|
||||
f" - {key}: [{content_type}] {value_preview} "
|
||||
f"(tags: {tags}, time: {timestamp[:19]})"
|
||||
)
|
||||
|
||||
if len(all_items) > 20:
|
||||
summary_lines.append(f" ... {len(all_items) - 20} more items not shown")
|
||||
|
||||
return "\n".join(summary_lines)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error while retrieving SketchPad summary: {str(e)}"
|
||||
|
||||
# Convenience methods for context management.
|
||||
def get_conversation_history(self, limit: Optional[int] = None):
|
||||
"""Get the conversation history for the current session."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return context.retrieve_messages(limit)
|
||||
|
||||
def get_full_saved_history(self, limit: Optional[int] = None):
|
||||
"""Get the fully saved conversation history."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return context.retrieve_full_messages(limit)
|
||||
|
||||
def search_conversation(self, query: str, limit: int = 5):
|
||||
"""Search the conversation history for the current session."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
# Use a simple search implementation.
|
||||
return context.search_messages(query, limit)
|
||||
|
||||
def search_full_history(self, query: str, limit: int = 5):
|
||||
"""Search the fully saved conversation history."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return context.search_messages(query, limit)
|
||||
|
||||
def clear_conversation(self) -> None:
|
||||
"""Clear the conversation history for the current session."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
context.clear_messages(keep_summary=True)
|
||||
|
||||
def get_conversation_summary(self) -> str:
|
||||
"""Get the conversation summary for the current session."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return context.get_summary() or ""
|
||||
|
||||
def get_full_saved_summary(self) -> str:
|
||||
"""Get the fully saved conversation summary."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return context.get_summary() or ""
|
||||
|
||||
def export_conversation(self, file_path: str) -> None:
|
||||
"""Export the conversation records for the current session."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
data = context.serialize()
|
||||
dir_path = os.path.dirname(file_path)
|
||||
if dir_path:
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def import_conversation(self, file_path: str, merge: bool = False) -> None:
|
||||
"""Import conversation records."""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if not merge:
|
||||
# Clear existing messages while preserving the summary.
|
||||
context.clear_messages(keep_summary=True)
|
||||
context.deserialize(data)
|
||||
|
||||
# Convenience methods for SketchPad management.
|
||||
async def store_in_sketch_pad(
|
||||
self,
|
||||
value,
|
||||
key: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
ttl: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Store data in SketchPad."""
|
||||
sketch_pad = get_current_sketch_pad()
|
||||
if sketch_pad is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
# Generate a key name if one is not provided.
|
||||
item_key = key or f"item_{uuid.uuid4().hex[:8]}"
|
||||
# Convert tags to a set.
|
||||
tags_set = set(tags) if tags else None
|
||||
await sketch_pad.set_item(
|
||||
key=item_key,
|
||||
value=value,
|
||||
ttl=ttl,
|
||||
summary=None,
|
||||
tags=tags_set,
|
||||
)
|
||||
return item_key
|
||||
|
||||
def get_from_sketch_pad(self, key: str) -> Any:
|
||||
"""Get data from SketchPad."""
|
||||
sketch_pad = get_current_sketch_pad()
|
||||
if sketch_pad is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return sketch_pad.get_value(key)
|
||||
|
||||
def search_sketch_pad(self, query: str, limit: int = 5):
|
||||
"""Search SketchPad content."""
|
||||
sketch_pad = get_current_sketch_pad()
|
||||
if sketch_pad is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return sketch_pad.search_by_content(query, limit)
|
||||
|
||||
def get_sketch_pad_stats(self):
|
||||
"""Get SketchPad statistics."""
|
||||
sketch_pad = get_current_sketch_pad()
|
||||
if sketch_pad is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
return sketch_pad.get_statistics()
|
||||
|
||||
def clear_sketch_pad(self):
|
||||
"""Clear SketchPad."""
|
||||
sketch_pad = get_current_sketch_pad()
|
||||
if sketch_pad is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
sketch_pad.clear()
|
||||
|
||||
def get_session_info(self):
|
||||
"""Get session information, including conversation history and SketchPad statistics."""
|
||||
try:
|
||||
conversation_count = len(self.get_conversation_history())
|
||||
sketch_pad_stats = self.get_sketch_pad_stats()
|
||||
conversation_summary = self.get_conversation_summary()
|
||||
except RuntimeError:
|
||||
# If there is no active conversation context, return basic information.
|
||||
conversation_count = 0
|
||||
sketch_pad_stats = {}
|
||||
conversation_summary = None
|
||||
|
||||
return {
|
||||
"agent_name": self.name,
|
||||
"model_name": self.model_name,
|
||||
"agent_class": self.__class__.__name__,
|
||||
"conversation_count": conversation_count,
|
||||
"sketch_pad_stats": sketch_pad_stats,
|
||||
"conversation_summary": conversation_summary,
|
||||
}
|
||||
|
||||
# ===== Common: streaming output and chronological persistence =====
|
||||
def _msg_to_dict(self, msg: Any) -> Dict[str, Any]:
|
||||
"""Convert backend-returned messages uniformly into dictionaries, supporting both object and dictionary forms."""
|
||||
if isinstance(msg, dict):
|
||||
return msg
|
||||
return {
|
||||
"role": getattr(msg, "role", None),
|
||||
"content": getattr(msg, "content", None),
|
||||
"tool_calls": getattr(msg, "tool_calls", None),
|
||||
"tool_call_id": getattr(msg, "tool_call_id", None),
|
||||
}
|
||||
|
||||
async def _stream_and_persist(
|
||||
self, response_packages: AsyncGenerator[Any, None]
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""
|
||||
Unified streaming processing and history persistence logic:
|
||||
- Continuously accumulate assistant text; when encountering tooluse/tool results, persist the accumulated text first, then write the tool message.
|
||||
- Ensure tool calls appear in history after the moment that triggered them, preserving the correct order.
|
||||
"""
|
||||
context = get_current_context()
|
||||
if context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
|
||||
assistant_buffer: str = ""
|
||||
baseline_len: Optional[int] = None
|
||||
|
||||
async for output in response_packages:
|
||||
yield output
|
||||
|
||||
if not is_response_yield(output):
|
||||
continue
|
||||
|
||||
current_messages = output.messages
|
||||
if baseline_len is None:
|
||||
try:
|
||||
baseline_len = (
|
||||
len(current_messages)
|
||||
if isinstance(current_messages, list)
|
||||
else 0
|
||||
)
|
||||
except Exception:
|
||||
baseline_len = 0
|
||||
|
||||
delta_text = extract_output_text(output, "agent_stream")
|
||||
if delta_text:
|
||||
assistant_buffer += delta_text
|
||||
|
||||
try:
|
||||
if isinstance(current_messages, list):
|
||||
curr_len = len(current_messages)
|
||||
if baseline_len is not None and curr_len > baseline_len:
|
||||
new_msgs = current_messages[baseline_len:curr_len]
|
||||
for nm in (self._msg_to_dict(x) for x in new_msgs):
|
||||
role = nm.get("role")
|
||||
content = nm.get("content")
|
||||
tool_calls = nm.get("tool_calls")
|
||||
tool_call_id = nm.get("tool_call_id")
|
||||
|
||||
if (role == "assistant" and tool_calls) or role == "tool":
|
||||
if assistant_buffer.strip():
|
||||
await context.store_message(
|
||||
Message(
|
||||
role="assistant", content=assistant_buffer
|
||||
)
|
||||
)
|
||||
assistant_buffer = ""
|
||||
|
||||
if role == "assistant" and tool_calls:
|
||||
await context.store_message(
|
||||
Message(
|
||||
role="assistant",
|
||||
content=None,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
)
|
||||
elif role == "tool":
|
||||
await context.store_message(
|
||||
Message(
|
||||
role="tool",
|
||||
content=content,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
)
|
||||
baseline_len = curr_len
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# At the end of the stream, write the remaining assistant text.
|
||||
if assistant_buffer.strip():
|
||||
await context.store_message(
|
||||
Message(role="assistant", content=assistant_buffer)
|
||||
)
|
||||
@@ -0,0 +1,345 @@
|
||||
from typing import Dict, List, Generator, Tuple, AsyncGenerator, override, Any
|
||||
from .BaseAgent import BaseAgent
|
||||
from context.conversation_manager import get_current_context
|
||||
from SimpleLLMFunc.type import ImgPath, ImgUrl, Text
|
||||
from tools import (
|
||||
execute_command,
|
||||
sketch_pad_operations,
|
||||
make_user_query_more_detailed,
|
||||
cad_code_generator,
|
||||
get_visual_feedback,
|
||||
)
|
||||
from context.schemas import Message
|
||||
|
||||
|
||||
class CADAgent(BaseAgent):
|
||||
@override
|
||||
def get_toolkit(self) -> List[Any]:
|
||||
return [
|
||||
make_user_query_more_detailed,
|
||||
cad_code_generator,
|
||||
execute_command,
|
||||
sketch_pad_operations,
|
||||
get_visual_feedback,
|
||||
]
|
||||
|
||||
@override
|
||||
def chat_impl(
|
||||
self,
|
||||
history: List[Dict[str, Any]],
|
||||
query: str | list[Text | ImgUrl | ImgPath | str],
|
||||
sketch_pad_summary: str,
|
||||
) -> Generator[Tuple[str, List[Dict[str, Any]]], None, None]: # type: ignore[override]
|
||||
"""
|
||||
# 🎯 Identity Description
|
||||
You are a professional intelligent assistant for CAD modeling, proficient in CADQuery/Python script-based modeling, geometric design, and engineering drawing. You communicate with the user in Chinese and are responsible for automating and closing the feedback loop across the full process from conceptual design to model export.
|
||||
|
||||
Your working directory contains or can access the SimpleCAD skill root reported in workspace facts; this folder is important because it holds documentation and operation guidance. When tools need API or workflow reference, they consult this preferred skill root.
|
||||
|
||||
---
|
||||
|
||||
# 🧭 Process Control Architecture (based on flowchart)
|
||||
|
||||
You follow the state machine architecture below to drive task execution. All state decisions and transitions are made autonomously by you.
|
||||
|
||||
## [State Nodes]
|
||||
|
||||
1. **Receive Requirement**
|
||||
- User inputs the modeling target
|
||||
- Use `make_user_query_more_detailed` tool to refine the input into a modeling specification, then store it as `req_xxxx`
|
||||
- Treat `make_user_query_more_detailed` as a specialist subagent: it can read SketchPad, inspect local skill docs directly, inspect APIs, and return a grounded requirement package
|
||||
- Do not assume the specialist should stop after a single failed tool call; it should continue searching/reading until it reaches an explicit terminal condition
|
||||
- In the `query` you pass to this tool, explicitly require the final detailed requirement to include exactly these headings:
|
||||
1. `## API Reference`
|
||||
2. `## Refined User Requirements`
|
||||
3. `## Parameter Table`
|
||||
4. `## Modeling Process`
|
||||
5. `## Notes`
|
||||
- Also explicitly require that the refined requirement be directly usable by `cad_code_generator`, including:
|
||||
- clear geometry goals,
|
||||
- dimensions and constraints,
|
||||
- assumptions and defaults,
|
||||
- reusable image paths and SketchPad ids,
|
||||
- verified API recommendations,
|
||||
- an ordered modeling sequence with recommended APIs and geometric intent.
|
||||
- Make it explicit that returning anything less than this structured requirement package counts as failure and the specialist should continue refining.
|
||||
- If the user uploaded reference images from Web UI, their workspace file paths may appear in the user message. When the image matters, pass the relevant path into `make_user_query_more_detailed(query_image_path=...)`.
|
||||
- Present the refined specification to the user and ask whether they are satisfied. If satisfied, proceed to **Confirm Requirement**;
|
||||
if not, return to **Receive Requirement**, integrate the new and old inputs, and refine again.
|
||||
|
||||
2. **Create Working Environment**
|
||||
- Decide the target folder and target script path: `./PartName_Specification/model.py`
|
||||
- If the folder does not exist yet, you may use `execute_command` to create it (`mkdir -p ...`)
|
||||
- Then proceed to **Generate Code**
|
||||
|
||||
3. **Generate Code**
|
||||
- Use `cad_code_generator` tool as a **single-call specialist subagent**
|
||||
- **CRITICAL**: You MUST pass `requirement_sketch_key` (the req_xxxx key from make_user_query_more_detailed). The specialist retrieves the detailed spec from SketchPad using this key.
|
||||
- Pass `task`, `target_file_path`, and `requirement_sketch_key`
|
||||
- `task` MUST explicitly include three parts: (1) mission, (2) context, (3) termination condition
|
||||
- Mission must clearly state: "Generate a SimpleCADAPI Python script to `target_file_path`"
|
||||
- Context must clearly state: user intent + `requirement_sketch_key` (which SketchPad key to read) + current failure/feedback (if any)
|
||||
- Termination condition must clearly state: "Run the script and export both STL and STEP/STP successfully; otherwise continue debugging"
|
||||
- Put all context into `task`: full user intent, create/modify type, modification target, traceback/visual feedback, success criteria
|
||||
- In `task`, explicitly instruct the specialist to validate and keep debugging until the script successfully exports the model files (normally STL and STEP/STP)
|
||||
- If API references are needed, tell the specialist to consult the preferred skill root's `references` directory briefly, then continue coding/debugging
|
||||
- In `task`, focus on objective/context/termination only; do not provide long workflow instructions
|
||||
- Do not assume the specialist should stop after one failed tool call or one failed patch; the intended behavior is to continue retrying until the explicit terminal condition is reached
|
||||
- Make it explicit that returning analysis without actually writing `target_file_path` counts as failure
|
||||
- The tool should be used when you need code creation, code modification, or self-contained code repair for the CAD script
|
||||
- Use this compact task template when calling `cad_code_generator`:
|
||||
`Mission: Generate a SimpleCADAPI script to <target_file_path>. Context: read SketchPad key <requirement_sketch_key>; user intent is <intent>; failure context is <errors/feedback>. Done only when the script has been executed and both <name>.stl and <name>.step/.stp are exported successfully.`
|
||||
|
||||
4. **Execute Modeling Script**
|
||||
- If `cad_code_generator` has just returned, first use `execute_command` to `ls` the target folder for `.stl` and `.step`/`.stp` outputs
|
||||
- If the expected exported files already exist, proceed directly to **Visual Feedback Check**
|
||||
- Otherwise use `execute_command` to run the script and export `.step`, `.stl` files
|
||||
- If successful, proceed to **Visual Feedback Check** and inform the user with artifact tags so downstream systems can recover the generated files reliably:
|
||||
```
|
||||
<|code_file|>Path to model.py</|code_file|>
|
||||
<|output_file|>Path to .stl file</|output_file|>
|
||||
<|output_file|>Path to .step or .stp file</|output_file|>
|
||||
```
|
||||
At minimum, always include the final `model.py` path with `<|code_file|>` and the final `.stl` path with `<|output_file|>`.
|
||||
Example:
|
||||
`<|code_file|>./Screwdriver_Slotted/model.py</|code_file|>`
|
||||
`<|output_file|>./Screwdriver_Slotted/Screwdriver.stl</|output_file|>`
|
||||
`<|output_file|>./Screwdriver_Slotted/Screwdriver.step</|output_file|>`
|
||||
- If execution fails or the model is not exported, enter Debug Phase (Traceback)
|
||||
|
||||
5. **Visual Feedback Check**
|
||||
- For visual verification, prefer passing the exported `.step`/`.stp` file into `get_visual_feedback`; use `.stl` only as a fallback when no CAD-native file is available
|
||||
- Use `render_multi_view_model` / `get_visual_feedback` to generate and inspect the 6-view image; output parameter must be the folder from **Create Working Environment**
|
||||
- If the result is unreasonable, automatically enter debug process; otherwise, go to **Task Completion**
|
||||
- Limit visual-feedback-based code repair to 3 rounds. After 3 failed visual feedback rounds, stop automatic edits, report current artifacts and remaining issues, and ask the user whether to continue.
|
||||
|
||||
6. **Task Completion**
|
||||
- Complete the modeling process and output a [Task Completed] message.
|
||||
- Repeat the final `<|code_file|>` and `<|output_file|>` tags in the completion message so artifact extraction remains stable.
|
||||
|
||||
---
|
||||
|
||||
## [Debug Subprocess: Traceback]
|
||||
|
||||
- Use execution error context to analyze the cause of failure
|
||||
- Pay attention to the printed structured information about solids during execution; use it to identify which edge/tag may require follow-up operations (e.g., fillet or chamfer)
|
||||
- Then use `cad_code_generator` with `requirement_sketch_key` (req_xxxx), `task`, and `target_file_path`.
|
||||
- **MUST pass requirement_sketch_key** so the specialist retrieves the detailed spec. **Always provide the full modeling goal every time you invoke it**
|
||||
- Prefer to send traceback / visual feedback / current intent back into `cad_code_generator` so the specialist subagent can repair the file itself
|
||||
- Debug responsibility is primarily on `cad_code_generator`; only handle debugging outside it when strictly necessary for workflow orchestration
|
||||
- When calling `cad_code_generator`, the `task` must clearly say whether it is a **modify-existing-code** task or a **create-new-file** task
|
||||
- In every debug call, `task` must still keep the same three-part structure: mission, context (`requirement_sketch_key` and latest failure), and termination condition (script executed and both STL + STEP/STP exported)
|
||||
- Pass enough context in `task` every time: full user intent, current target file path, what changed, what failed, and what success should look like
|
||||
- Also make `task` explicitly require validation and debugging until STL and STEP/STP export succeeds; a clean script run without exported models is **not** enough
|
||||
- If the specialist needs API docs, point it to the preferred skill root's `references` directory
|
||||
- Assume the specialist runs under the workspace root, so all paths in `task` should be workspace-relative or absolute and should not reference repo-only helper paths
|
||||
- When you regain control after `cad_code_generator`, prefer `ls` export checks before rerunning the script; rerun only when outputs are missing or a fresh traceback is needed
|
||||
- Return to Step 4 **Execute Modeling Script** to retry after applying necessary fixes (feedback, traceback, etc.), until model export and render are successful; however, visual-feedback-based repair must stop after 3 rounds and wait for user confirmation.
|
||||
|
||||
---
|
||||
|
||||
# 🧰 Tool Usage Guide (Summary)
|
||||
|
||||
| Tool Name | Purpose |
|
||||
|-----------|---------|
|
||||
| `make_user_query_more_detailed` | Specialist subagent for refining vague requirements into a grounded modeling spec; may inspect docs/APIs/SketchPad before saving `req_xxxx` |
|
||||
| `cad_code_generator` | Generate or modify SimpleCADAPI code based on requirement and write directly to target file |
|
||||
| `execute_command` | Execute Python scripts or shell commands |
|
||||
| `sketch_pad_operations` | Perform `store/retrieve/search/delete/list/search_tags/search/clear/stats` on SketchPad |
|
||||
| `get_visual_feedback` | Given the requirement, code, and model path, return targeted visual feedback and modification suggestions. Prefer STEP/STP as `model_path` when available |
|
||||
|
||||
Notes on `cad_code_generator`:
|
||||
- Treat it as a specialist coding subagent, not as a dumb text generator.
|
||||
- **MUST pass requirement_sketch_key** (req_xxxx from make_user_query_more_detailed). The specialist uses it to retrieve the detailed spec.
|
||||
- Pass `task`, `target_file_path`, and `requirement_sketch_key`.
|
||||
- Your responsibility is to pack the right context into `task` and always provide the requirement key.
|
||||
- Its responsibility is to create/repair `model.py` until validation passes and the model export succeeds, or it can clearly explain what extra context is still missing.
|
||||
|
||||
---
|
||||
|
||||
# 📂 File Organization Requirements
|
||||
|
||||
- Folder structure: `./PartType_Parameters/`
|
||||
- Script filename: `model.py`
|
||||
- Execute via: `cd ./PartType_Parameters && uv run python model.py` to ensure outputs reside in the same folder
|
||||
- Export check via: `cd ./PartType_Parameters && ls *.stl && (ls *.step || ls *.stp)`
|
||||
- Output file names: `PartName.step`, `PartName.stl`
|
||||
- Render image name: `PartName_multi_view_render.png`
|
||||
- All files must be in the same folder
|
||||
- When you know the generated script path, include `<|code_file|>...</|code_file|>` in an assistant message using a workspace-relative path when possible
|
||||
- When you know generated model artifact paths, include `<|output_file|>...</|output_file|>` for each key artifact you want surfaced; this must include the final `.stl`, and should include the final `.step`/`.stp` when available
|
||||
|
||||
---
|
||||
|
||||
# 🧠 Intelligent Behavior Constraints
|
||||
|
||||
- Only store data in SketchPad when it is truly needed for cross-step context; do not save generated code into SketchPad by default
|
||||
- After a failure, return to the corresponding state and attempt to fix it
|
||||
- Do not proceed past user confirmation in any phase
|
||||
- If the task is too complex, proactively suggest requirement decomposition
|
||||
|
||||
---
|
||||
|
||||
# 💡 Important Notes
|
||||
|
||||
- Always call `sketch_pad_operations: clear` before starting each task to reset the environment
|
||||
- Do not generate code or write files without user confirmation
|
||||
- Visual verification stage is mandatory; must use `get_visual_feedback` and rigorously follow suggestions
|
||||
- Visual-feedback-based automatic repair is capped at 3 rounds per task; execution/export failures may be retried more, but repeated visual mismatches require user confirmation after the cap.
|
||||
- When both STL and STEP/STP exist, visual feedback should use the STEP/STP file because it renders CAD geometry more faithfully
|
||||
- Once user confirms satisfaction, summarize the modeling process and declare "This modeling task is completed, code archived".
|
||||
- Clearly indicate the **current state node**, and state clearly what will be done next at the end of each state
|
||||
- Tool calls must be clearly announced before use with appropriate emoji
|
||||
|
||||
Example Response:
|
||||
### Entering [Receive Requirement] State
|
||||
|
||||
- We will use the `make_user_query_more_detailed` tool to convert vague input into a structured modeling spec.
|
||||
|
||||
I now know the detailed modeling spec is saved in SketchPad under key: xxxxxx
|
||||
|
||||
### Entering [Confirm Requirement] State
|
||||
|
||||
- Next, I will use the `sketch_pad_operations` tool to retrieve and display the spec for your confirmation.
|
||||
|
||||
#### Modeling Requirement:
|
||||
|
||||
xxxxxxxx
|
||||
|
||||
Does this description match your expectations? Would you like any changes? If confirmed, we will proceed to [Create Working Environment]; otherwise, we will iterate again until you are satisfied.
|
||||
|
||||
- Always use a line break for state entry descriptions and mark the current state with `###` headings.
|
||||
- Use proper emoji for tool invocation.
|
||||
- When using file operation tools to debug, ensure correct indentation. Read context before writing.
|
||||
- If reusing `cad_code_generator`, always provide full user intent + API reference + current code.
|
||||
- Always converse with user in English.
|
||||
"""
|
||||
return # type: ignore[return-value]
|
||||
|
||||
async def run(
|
||||
self, query: Any, raw_user_content: Any = None
|
||||
) -> AsyncGenerator[Any, None]: # type: ignore[override]
|
||||
"""Run the agent with the given query.
|
||||
|
||||
Args:
|
||||
query (str): The query to process.
|
||||
|
||||
Returns:
|
||||
Generator[str, None, None]: The response chunks from the agent.
|
||||
"""
|
||||
if not query:
|
||||
raise ValueError("Query must not be empty")
|
||||
|
||||
# Get the SketchPad keys and truncated value contents.
|
||||
sketch_pad_summary = self.get_sketch_pad_summary()
|
||||
|
||||
import re
|
||||
|
||||
raw_query = raw_user_content if raw_user_content is not None else query
|
||||
|
||||
def _normalize_query_for_chat(content: Any) -> Any:
|
||||
if isinstance(content, str):
|
||||
query_image_match = re.search(r"<i>(.*?)</i>", content)
|
||||
if query_image_match:
|
||||
query_image = query_image_match.group(1)
|
||||
text_query = re.sub(r"<i>.*?</i>", "", content).strip()
|
||||
try:
|
||||
return (
|
||||
[Text(text_query), ImgPath(query_image)]
|
||||
if text_query
|
||||
else [ImgPath(query_image)]
|
||||
)
|
||||
except Exception:
|
||||
if text_query:
|
||||
return (
|
||||
text_query
|
||||
+ f" (Attached image: the image file path is {query_image}. You can pass this path to `make_user_query_more_detailed` or later to the `get_visual_feedback` tool to obtain more information.)"
|
||||
)
|
||||
return re.sub(r"<i>.*?</i>", "", content).strip()
|
||||
|
||||
if isinstance(content, list):
|
||||
normalized_parts: list[Text | ImgUrl | ImgPath] = []
|
||||
for item in content:
|
||||
if isinstance(item, Text):
|
||||
normalized_parts.append(item)
|
||||
continue
|
||||
if isinstance(item, ImgUrl):
|
||||
normalized_parts.append(item)
|
||||
continue
|
||||
if isinstance(item, ImgPath):
|
||||
normalized_parts.append(item)
|
||||
continue
|
||||
if isinstance(item, str):
|
||||
if item.strip():
|
||||
normalized_parts.append(Text(item))
|
||||
continue
|
||||
|
||||
item_type = getattr(item, "type", None) or (
|
||||
item.get("type") if isinstance(item, dict) else None
|
||||
)
|
||||
if item_type == "text":
|
||||
text_value = getattr(item, "text", None) or (
|
||||
item.get("text") if isinstance(item, dict) else None
|
||||
)
|
||||
if isinstance(text_value, str) and text_value.strip():
|
||||
normalized_parts.append(Text(text_value))
|
||||
elif item_type == "image_url":
|
||||
image_payload = getattr(item, "image_url", None) or (
|
||||
item.get("image_url") if isinstance(item, dict) else None
|
||||
)
|
||||
if isinstance(image_payload, dict) and isinstance(
|
||||
image_payload.get("url"), str
|
||||
):
|
||||
detail = image_payload.get("detail")
|
||||
if not isinstance(detail, str):
|
||||
detail = "auto"
|
||||
normalized_parts.append(
|
||||
ImgUrl(image_payload["url"], detail=detail)
|
||||
)
|
||||
|
||||
if len(normalized_parts) == 1 and isinstance(normalized_parts[0], Text):
|
||||
return str(normalized_parts[0])
|
||||
return normalized_parts
|
||||
|
||||
return str(content)
|
||||
|
||||
chat_query = _normalize_query_for_chat(query)
|
||||
|
||||
# Get the current conversation context.
|
||||
current_context = get_current_context()
|
||||
if current_context is None:
|
||||
raise RuntimeError("No active conversation context")
|
||||
|
||||
def _normalize_message_content_for_history(content: Any) -> Any:
|
||||
if isinstance(content, str) or content is None:
|
||||
return content or ""
|
||||
if isinstance(content, list):
|
||||
normalized_items: List[Dict[str, Any]] = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
normalized_items.append(item)
|
||||
elif hasattr(item, "model_dump"):
|
||||
normalized_items.append(item.model_dump())
|
||||
return normalized_items
|
||||
return str(content)
|
||||
|
||||
history_messages = current_context.retrieve_messages()
|
||||
history: List[Dict[str, Any]] = []
|
||||
for m in history_messages:
|
||||
if m.role in ("user", "assistant"):
|
||||
history.append(
|
||||
{
|
||||
"role": m.role,
|
||||
"content": _normalize_message_content_for_history(m.content),
|
||||
}
|
||||
)
|
||||
|
||||
# Before starting the conversation, store the current user message in the context.
|
||||
await current_context.store_message(Message(role="user", content=raw_query))
|
||||
|
||||
# Call the LLM in raw streaming mode.
|
||||
response_packages = self.chat(history, chat_query, sketch_pad_summary)
|
||||
|
||||
# Reuse the base class streaming processing and history persistence logic.
|
||||
async for raw in self._stream_and_persist(response_packages):
|
||||
yield raw
|
||||
@@ -0,0 +1,25 @@
|
||||
from .BaseAgent import BaseAgent
|
||||
from .AgentRegister import (
|
||||
AgentRegistry,
|
||||
get_agent_registry,
|
||||
register_agent,
|
||||
get_agent,
|
||||
list_available_models
|
||||
)
|
||||
|
||||
|
||||
from .CADAgent import CADAgent
|
||||
|
||||
|
||||
__all__ = [
|
||||
'BaseAgent',
|
||||
'AgentRegistry',
|
||||
'get_agent_registry',
|
||||
'register_agent',
|
||||
'get_agent',
|
||||
'list_available_models',
|
||||
]
|
||||
|
||||
|
||||
# Register CADAgent.
|
||||
register_agent("cadagent", CADAgent)
|
||||
Reference in New Issue
Block a user