adding default Basemodel

pull/648/head
Occupying-Mars 5 months ago
parent 47a359ec34
commit 019105f154

@ -76,6 +76,38 @@ def summarize_expenses(data: SummarizeExpenses) -> dict:
return {"total_expenses": total_expenses}
```
#### Using Default BaseModel
When initializing an Agent with tools but without specifying `list_base_models`, the system automatically uses a DefaultBaseModel:
```python
from pydantic import BaseModel, Field
from typing import Any
class DefaultBaseModel(BaseModel):
"""A default base model that accepts any data structure"""
data: Any = Field(default=None, description="Any data type")
class Config:
arbitrary_types_allowed = True
extra = "allow"
# Agent will use DefaultBaseModel automatically if no list_base_models provided
agent = Agent(
agent_name="Simple Agent",
tools=[your_tool_function], # DefaultBaseModel will be used
llm=your_llm_instance
)
```
⚠️ **Warning**: While the DefaultBaseModel provides flexibility, it's recommended to define your own specific Pydantic models for:
- Type safety and validation
- Clear documentation
- Predictable behavior
- Schema generation
For production use, define custom models like the examples above (CalculateTax, GenerateInvoice).
#### Using Functions Directly
Tools can also be defined directly as functions without using Pydantic models. This approach is suitable for simpler tasks where complex validation is not required.

@ -26,7 +26,7 @@ from clusterops import (
execute_on_gpu,
execute_with_cpu_cores,
)
from pydantic import BaseModel
from pydantic import BaseModel, Field
from swarm_models.tiktoken_wrapper import TikTokenizer
from termcolor import colored
@ -90,6 +90,14 @@ ToolUsageType = Union[BaseModel, Dict[str, Any]]
# [FEAT][AGENT]
class DefaultBaseModel(BaseModel):
"""A default base model that accepts any data structure"""
data: Any = Field(default=None, description="Any data type")
class Config:
arbitrary_types_allowed = True
extra = "allow"
class Agent:
"""
Agent is the backbone to connect LLMs with tools and long term memory. Agent also provides the ability to
@ -477,6 +485,24 @@ class Agent:
or exists(tool_schema)
):
# If no base models provided, use the default one
if exists(tools) or exists(tool_schema):
if list_base_models is None:
list_base_models = [DefaultBaseModel]
logger.warning(
colored(
"\nUsing default base model for tool schema validation. "
"This can lead to potential format mismatches or validation issues. "
"Consider defining your own Pydantic BaseModel(s) with proper field definitions "
"and validation rules for your specific use case.\n"
"Example:\n"
"class YourModel(BaseModel):\n"
" field1: str = Field(..., description='Description of field1')\n"
" field2: int = Field(..., description='Description of field2')\n",
"yellow"
)
)
self.tool_struct = BaseTool(
tools=tools,
base_models=list_base_models,

Loading…
Cancel
Save