From 4526ac3a367a527fc8cc6a4161871f90566a453e Mon Sep 17 00:00:00 2001
From: Nicholas Fleischhauer <unobtainiumrock@gmail.com>
Date: Mon, 10 Feb 2025 12:47:46 -0800
Subject: [PATCH] Migrated out pydantic models to keep api.py clean and focused
 primarily on API logic. Pydantic models will now be located under
 /pydantic_models nad imported as from pydantic_models import *

---
 api.py                      | 23 +---------------------
 pydantic_models/__init__.py | 13 +++++++++++++
 pydantic_models/models.py   | 38 +++++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 22 deletions(-)
 create mode 100644 pydantic_models/__init__.py
 create mode 100644 pydantic_models/models.py

diff --git a/api.py b/api.py
index ba82edf..aeeb09b 100644
--- a/api.py
+++ b/api.py
@@ -1,12 +1,11 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 import adalflow as adal
 from src.rag import RAG
-from typing import List, Optional
 import os
 from dotenv import load_dotenv
 from datetime import datetime, timezone
+from pydantic_models import *
 
 def load_environment():
     """Load environment variables from .env file if available, otherwise use system environment variables."""
@@ -59,26 +58,6 @@ def load_environment():
     print(f"Error initializing RAG component: {e}")
     raise RuntimeError(f"Failed to initialize RAG component: {e}")
 
-class QueryRequest(BaseModel):
-    repo_url: str
-    query: str
-
-class DocumentMetadata(BaseModel):
-    file_path: str
-    type: str
-    is_code: bool = False
-    is_implementation: bool = False
-    title: str = ""
-
-class Document(BaseModel):
-    text: str
-    meta_data: DocumentMetadata
-
-class QueryResponse(BaseModel):
-    rationale: str
-    answer: str
-    contexts: List[Document]
-
 @app.post("/query", response_model=QueryResponse)
 async def query_repository(request: QueryRequest):
     """
diff --git a/pydantic_models/__init__.py b/pydantic_models/__init__.py
new file mode 100644
index 0000000..98727a7
--- /dev/null
+++ b/pydantic_models/__init__.py
@@ -0,0 +1,13 @@
+from .models import (
+    QueryRequest,
+    QueryResponse,
+    Document,
+    DocumentMetadata
+)
+
+__all__ = [
+    "QueryRequest",
+    "QueryResponse",
+    "Document",
+    "DocumentMetadata"
+]
\ No newline at end of file
diff --git a/pydantic_models/models.py b/pydantic_models/models.py
new file mode 100644
index 0000000..6dd1e18
--- /dev/null
+++ b/pydantic_models/models.py
@@ -0,0 +1,38 @@
+from pydantic import BaseModel
+from typing import List
+
+class QueryRequest(BaseModel):
+    repo_url: str
+    query: str
+
+class DocumentMetadata(BaseModel):
+    file_path: str
+    type: str
+    is_code: bool = False
+    is_implementation: bool = False
+    title: str = ""
+
+# Consider extending what metadata is tracked in DoumentMetaData
+# or create an additional class. I think it could be important to track other
+# metadata metrics, such as line count and file size.
+
+# Justification: 
+# During the EDA stages of model improvement, tracking repo sizes, file sizes, line counts
+# could reveal relationships between performance.
+
+class Document(BaseModel):
+    text: str
+    meta_data: DocumentMetadata
+
+class QueryResponse(BaseModel):
+    rationale: str
+    answer: str
+    contexts: List[Document]
+
+
+__all__ = [
+    "QueryReqest",
+    "QueryResponse",
+    "Document",
+    "DocumentMetadata"
+]
\ No newline at end of file