mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
from typing import Any, Optional, Type
|
|
|
|
import instructor
|
|
from pydantic import BaseModel, Field, PrivateAttr, model_validator
|
|
|
|
|
|
class Instructor(BaseModel):
|
|
"""Class that wraps an agent llm with instructor."""
|
|
|
|
_client: Any = PrivateAttr()
|
|
content: str = Field(description="Content to be sent to the instructor.")
|
|
agent: Optional[Any] = Field(
|
|
description="The agent that needs to use instructor.", default=None
|
|
)
|
|
llm: Optional[Any] = Field(
|
|
description="The agent that needs to use instructor.", default=None
|
|
)
|
|
instructions: Optional[str] = Field(
|
|
description="Instructions to be sent to the instructor.",
|
|
default=None,
|
|
)
|
|
model: Type[BaseModel] = Field(
|
|
description="Pydantic model to be used to create an output."
|
|
)
|
|
|
|
@model_validator(mode="after")
|
|
def set_instructor(self):
|
|
"""Set instructor."""
|
|
if self.agent and not self.llm:
|
|
self.llm = self.agent.function_calling_llm or self.agent.llm
|
|
|
|
self._client = instructor.patch(
|
|
self.llm.client._client,
|
|
mode=instructor.Mode.TOOLS,
|
|
)
|
|
return self
|
|
|
|
def to_json(self):
|
|
model = self.to_pydantic()
|
|
return model.model_dump_json(indent=2)
|
|
|
|
def to_pydantic(self):
|
|
messages = [{"role": "user", "content": self.content}]
|
|
if self.instructions:
|
|
messages.append({"role": "system", "content": self.instructions})
|
|
|
|
model = self._client.chat.completions.create(
|
|
model=self.llm.model_name, response_model=self.model, messages=messages
|
|
)
|
|
return model
|