mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import warnings
|
|
from typing import Any, Optional, Type
|
|
|
|
|
|
class InternalInstructor:
|
|
"""Class that wraps an agent llm with instructor."""
|
|
|
|
def __init__(
|
|
self,
|
|
content: str,
|
|
model: Type,
|
|
agent: Optional[Any] = None,
|
|
llm: Optional[str] = None,
|
|
):
|
|
self.content = content
|
|
self.agent = agent
|
|
self.llm = llm
|
|
self.model = model
|
|
self._client = None
|
|
self.set_instructor()
|
|
|
|
def set_instructor(self):
|
|
"""Set instructor."""
|
|
if self.agent and not self.llm:
|
|
self.llm = self.agent.function_calling_llm or self.agent.llm
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore", UserWarning)
|
|
import instructor
|
|
from litellm import completion
|
|
|
|
is_custom_openai = getattr(self.llm, 'model', '').startswith('custom_openai/')
|
|
|
|
mode = instructor.Mode.PARALLEL_TOOLS if is_custom_openai else instructor.Mode.TOOLS
|
|
|
|
self._client = instructor.from_litellm(
|
|
completion,
|
|
mode=mode,
|
|
)
|
|
|
|
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}]
|
|
model = self._client.chat.completions.create(
|
|
model=self.llm.model, response_model=self.model, messages=messages
|
|
)
|
|
|
|
if isinstance(model, list) and len(model) > 0:
|
|
return model[0] # Return the first model from the list
|
|
|
|
return model
|