refactor: new module for pydantic models

app/main.py

@@ -1,34 +1,18 @@
 from fastapi import FastAPI
 from google import genai
-from pydantic import BaseModel
 
-from app import agent, config
+from app import agent, config, models
 
 gemini = genai.Client(api_key=config.settings.google_api_key)
 app = FastAPI()
 
 
-
-class Source(BaseModel):
-    title: str
-    url: str
-
-
-class Statement(BaseModel):
-    text: str
-    sources: list[Source] | None = None
-
-
-class Statements(BaseModel):
-    statements: list[Statement]
-
-
 @app.get("/health")
 def health_check():
     return "ok"
 
 
-@app.get("/ask", response_model=Statements)
+@app.get("/ask", response_model=models.Statements)
 def ask(query: str):
     output = agent.respond(gemini, query)
     return {"statements": output}

app/models.py

@@ -0,0 +1,27 @@
+from pydantic import BaseModel, computed_field
+
+
+class Source(BaseModel):
+    title: str
+    url: str
+    summary: str | None = None
+    abstract: str | None = None
+    doi: str | None = None
+    venue: str | None = None
+    year: int | None = None
+    citations: int | None = None
+    influential_citations: int | None = None
+
+    @computed_field
+    @property
+    def citation(self) -> str:
+        return f"[{self.title}]({self.url}). *{self.venue}* {self.year}."
+
+
+class Statement(BaseModel):
+    text: str
+    sources: list[Source] | None = None
+
+
+class Statements(BaseModel):
+    statements: list[Statement]