Implemented pydantic

api_clients.py

@@ -4,6 +4,7 @@ import aiohttp
 import elevenlabs
 from contextlib import asynccontextmanager
 from logger import setup_logger, log_execution_time, log_async_execution_time
+from models import OpenRouterRequest, OpenRouterResponse, Message, OpenRouterModel
 
 logger = setup_logger("api_clients")
 
@@ -60,9 +61,8 @@ class OpenRouterClient:
             headers=self.headers
         )
         response.raise_for_status()
-        models = response.json()["data"]
-        # Return tuple of (id, display_name) for each model
-        return [(model["id"], model["name"]) for model in models]
+        models = [OpenRouterModel(**model) for model in response.json()["data"]]
+        return [(model.id, model.name) for model in models]
 
     @log_async_execution_time(logger)
     async def generate_script(self, content: str, prompt: str, model_id: str) -> str:
@@ -84,19 +84,23 @@ class OpenRouterClient:
             raise ValueError("Please provide a more detailed prompt")
         
         try:
+            request = OpenRouterRequest(
+                model=model_id,
+                messages=[
+                    Message(role="system", content="You are a podcast script writer."),
+                    Message(role="user", content=f"Create a podcast script from this content: {content}")
+                ]
+            )
+            
             response = requests.post(
                 f"{self.base_url}/chat/completions",
                 headers=self.headers,
-                json={
-                    "model": model_id,
-                    "messages": [
-                        {"role": "system", "content": "You are a podcast script writer."},
-                        {"role": "user", "content": f"Create a podcast script from this content: {content}"}
-                    ]
-                }
+                data=request.json()
             )
             response.raise_for_status()
-            return response.json()["choices"][0]["message"]["content"]
+            
+            router_response = OpenRouterResponse(**response.json())
+            return router_response.choices[0].message.content
         except Exception as e:
             logger.error(f"Script generation failed", exc_info=True)
             raise

models.py

@@ -0,0 +1,25 @@
+from pydantic import BaseModel, Field
+from typing import List, Optional
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+class OpenRouterRequest(BaseModel):
+    model: str
+    messages: List[Message]
+
+class Choice(BaseModel):
+    message: Message
+    index: int = 0
+    finish_reason: Optional[str] = None
+
+class OpenRouterResponse(BaseModel):
+    id: str
+    choices: List[Choice]
+    model: str
+
+class OpenRouterModel(BaseModel):
+    id: str
+    name: str
+    description: Optional[str] = None