Use pydantic ai

Author: YanickSchraner

README.md

@@ -1,3 +1,8 @@
+# TODO
+returns library
+api error from backend common
+
+
 # Transcribo Backend
 
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "dcc-backend-common[fastapi]>=0.0.2",
     "dependency-injector>=4.48.3",
     "fastapi[standard]>=0.128.0",
+    "pydantic-ai>=1.43.0",
     "returns>=0.26.0",
 ]
 

src/transcribo_backend/agents/__init__.py

@@ -0,0 +1,34 @@
+from openai import AsyncOpenAI
+from pydantic_ai import Agent, TextOutput
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from transcribo_backend.utils.app_config import AppConfig
+
+
+def transform_to_swissgerman_style(text: str) -> str:
+    return text.replace("ß", "ss")
+
+
+def create_summarize_agent(app_config: AppConfig) -> Agent:
+    client = AsyncOpenAI(max_retries=3, base_url=app_config.llm_base_url, api_key=app_config.api_key)
+    model = OpenAIChatModel(model_name=app_config.llm_model, provider=OpenAIProvider(openai_client=client))
+    summarize_agent: Agent = Agent(
+        model=model,
+        output_type=TextOutput(transform_to_swissgerman_style),
+    )
+
+    @summarize_agent.instructions
+    def get_instructions() -> str:
+        return """
+You are a meeting summary expert.
+You are given a transcript of a meeting and you need to summarize it.
+You need to summarize the meeting in a way that is easy to understand and use.
+You need to include the main points of the meeting, the decisions made, and the action items.
+You need to include the names of the participants.
+You use markdown to format the summary.
+Use the same language as used in the transcript to summarize the meeting.
+If you are not sure about the language, use German.
+"""
+
+    return summarize_agent

src/transcribo_backend/agents/summarize_agent.py

@@ -1,44 +1,19 @@
-from openai import AsyncOpenAI
-
+from transcribo_backend.agents.summarize_agent import create_summarize_agent
 from transcribo_backend.models.summary import Summary
 from transcribo_backend.utils.app_config import AppConfig
 
-system_prompt = """
-You are a meeting summary expert.
-You are given a transcript of a meeting and you need to summarize it.
-You need to summarize the meeting in a way that is easy to understand and use.
-You need to include the main points of the meeting, the decisions made, and the action items.
-You need to include the names of the participants.
-You use markdown to format the summary.
-Use the same language as used in the transcript to summarize the meeting.
-If you are not sure about the language, use German.
-"""
-
 
 class SummarizationService:
     def __init__(self, app_config: AppConfig):
         self.app_config = app_config
+        self.agent = create_summarize_agent(app_config)
 
     async def summarize(self, transcript: str) -> Summary:
         """
         Summarize a transcript of a meeting.
         """
         try:
-            client = AsyncOpenAI(api_key=self.app_config.api_key, base_url=self.app_config.api_key)
-            models = await client.models.list()
-
-            if not models.data:
-                raise ValueError("Could not find any models available from the API")
-
-            model = models.data[0].id
-        except Exception as e:
-            raise RuntimeError("Failed to initialize OpenAI client or fetch models.") from e
-
-        try:
-            response = await client.chat.completions.create(
-                model=model,
-                messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": transcript}],
-            )
-            return Summary(summary=response.choices[0].message.content)
+            result = await self.agent.run(transcript)
+            return Summary(summary=result.output)
         except Exception as e:
             raise RuntimeError("Failed to generate summary.") from e

src/transcribo_backend/services/summarization_service.py

@@ -0,0 +1,38 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from transcribo_backend.models.summary import Summary
+from transcribo_backend.services.summarization_service import SummarizationService
+from transcribo_backend.utils.app_config import AppConfig
+
+
+@pytest.mark.asyncio
+async def test_summarize_calls_agent():
+    # Mock AppConfig
+    app_config = MagicMock(spec=AppConfig)
+
+    # Mock the Agent and its run method
+    mock_agent = MagicMock()
+    mock_run_result = MagicMock()
+    mock_run_result.data = "This is a summary."
+    mock_agent.run = AsyncMock(return_value=mock_run_result)
+
+    # Patch create_summarize_agent to return our mock agent
+    with patch(
+        "transcribo_backend.services.summarization_service.create_summarize_agent", return_value=mock_agent
+    ) as mock_create:
+        service = SummarizationService(app_config)
+
+        # Verify agent creation was called
+        mock_create.assert_called_once_with(app_config)
+
+        transcript = "Some meeting transcript."
+        result = await service.summarize(transcript)
+
+        # Verify run was called
+        mock_agent.run.assert_called_once_with(transcript)
+
+        # Verify result
+        assert isinstance(result, Summary)
+        assert result.summary == "This is a summary."