KeyArgo
/
smolagents
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: update model type to ChatMessage in agent classes (#263)

This commit is contained in:
Jan 2025-01-18 18:27:48 +01:00 committed by GitHub
parent e3ae1800e2
commit 34810986e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/smolagents/agents.py
View File

@ -34,7 +34,10 @@ from .local_python_executor import (
LocalPythonInterpreter, LocalPythonInterpreter,
fix_final_answer_code, fix_final_answer_code,
) )
from .models import MessageRole from .models import (
MessageRole,
ChatMessage,
)
from .monitoring import Monitor from .monitoring import Monitor
from .prompts import ( from .prompts import (
CODE_SYSTEM_PROMPT, CODE_SYSTEM_PROMPT,
@ -191,7 +194,7 @@ class MultiStepAgent:
def __init__( def __init__(
self, self,
tools: List[Tool], tools: List[Tool],
model: Callable[[List[Dict[str, str]]], str], model: Callable[[List[Dict[str, str]]], ChatMessage],
system_prompt: Optional[str] = None, system_prompt: Optional[str] = None,
tool_description_template: Optional[str] = None, tool_description_template: Optional[str] = None,
max_steps: int = 6, max_steps: int = 6,
@ -775,7 +778,7 @@ class ToolCallingAgent(MultiStepAgent):
def __init__( def __init__(
self, self,
tools: List[Tool], tools: List[Tool],
model: Callable, model: Callable[[List[Dict[str, str]]], ChatMessage],
system_prompt: Optional[str] = None, system_prompt: Optional[str] = None,
planning_interval: Optional[int] = None, planning_interval: Optional[int] = None,
**kwargs, **kwargs,
@ -885,7 +888,7 @@ class CodeAgent(MultiStepAgent):
def __init__( def __init__(
self, self,
tools: List[Tool], tools: List[Tool],
model: Callable, model: Callable[[List[Dict[str, str]]], ChatMessage],
system_prompt: Optional[str] = None, system_prompt: Optional[str] = None,
grammar: Optional[Dict[str, str]] = None, grammar: Optional[Dict[str, str]] = None,
additional_authorized_imports: Optional[List[str]] = None, additional_authorized_imports: Optional[List[str]] = None,