Update OpenAI pydantic model

This commit is contained in:
Yiorgis Gozadinos 2025-09-03 11:41:48 +03:00
parent 9830726bc7
commit 450a58eb52
No known key found for this signature in database

View file

@ -1,6 +1,6 @@
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIModel from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.ollama import OllamaProvider from pydantic_ai.providers.ollama import OllamaProvider
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
@ -61,7 +61,7 @@ class QuestionAnswerAgent:
def _get_model(self, provider: str, model: str): def _get_model(self, provider: str, model: str):
"""Get the appropriate model object for the provider.""" """Get the appropriate model object for the provider."""
if provider == "ollama": if provider == "ollama":
return OpenAIModel( return OpenAIChatModel(
model_name=model, model_name=model,
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
) )