Update main.py

main.py

@@ -1,6 +1,6 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query, Path
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from typing import List
 import json
 import os
@@ -11,7 +11,11 @@ from txtai.embeddings import Embeddings
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
+app = FastAPI(
+    title="Embeddings API",
+    description="An API for creating and querying text embeddings indexes.",
+    version="1.0.0"
+)
 
 # Enable CORS
 app.add_middleware(
@@ -25,22 +29,20 @@ app.add_middleware(
 embeddings = Embeddings({"path": "avsolatorio/GIST-all-MiniLM-L6-v2"})
 
 class DocumentRequest(BaseModel):
-    index_id: str
-    documents: List[str]
+    index_id: str = Field(..., description="Unique identifier for the index")
+    documents: List[str] = Field(..., description="List of documents to be indexed")
 
 class QueryRequest(BaseModel):
-    index_id: str
-    query: str
-    num_results: int
+    index_id: str = Field(..., description="Unique identifier for the index to query")
+    query: str = Field(..., description="The search query")
+    num_results: int = Field(..., description="Number of results to return", ge=1)
 
-def save_embeddings(index_id, document_list):
+def save_embeddings(index_id: str, document_list: List[str]):
     try:
         folder_path = f"/app/indexes/{index_id}"
         os.makedirs(folder_path, exist_ok=True)
-        
         # Save embeddings
         embeddings.save(f"{folder_path}/embeddings")
-        
         # Save document_list
         with open(f"{folder_path}/document_list.json", "w") as f:
             json.dump(document_list, f)
@@ -49,29 +51,31 @@ def save_embeddings(index_id, document_list):
         logger.error(f"Error saving embeddings for index_id {index_id}: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error saving embeddings: {str(e)}")
 
-def load_embeddings(index_id):
+def load_embeddings(index_id: str) -> List[str]:
     try:
         folder_path = f"/app/indexes/{index_id}"
-        
         if not os.path.exists(folder_path):
             logger.error(f"Index not found for index_id: {index_id}")
             raise HTTPException(status_code=404, detail="Index not found")
-        
         # Load embeddings
         embeddings.load(f"{folder_path}/embeddings")
-        
         # Load document_list
         with open(f"{folder_path}/document_list.json", "r") as f:
             document_list = json.load(f)
-        
         logger.info(f"Embeddings and document list loaded for index_id: {index_id}")
         return document_list
     except Exception as e:
         logger.error(f"Error loading embeddings for index_id {index_id}: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error loading embeddings: {str(e)}")
 
-@app.post("/create_index/")
+@app.post("/create_index/", response_model=dict, tags=["Index Operations"])
 async def create_index(request: DocumentRequest):
+    """
+    Create a new index with the given documents.
+    
+    - **index_id**: Unique identifier for the index
+    - **documents**: List of documents to be indexed
+    """
     try:
         document_list = [(i, text, None) for i, text in enumerate(request.documents)]
         embeddings.index(document_list)
@@ -82,8 +86,15 @@ async def create_index(request: DocumentRequest):
         logger.error(f"Error creating index: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error creating index: {str(e)}")
 
-@app.post("/query_index/")
+@app.post("/query_index/", response_model=dict, tags=["Index Operations"])
 async def query_index(request: QueryRequest):
+    """
+    Query an existing index with the given search query.
+    
+    - **index_id**: Unique identifier for the index to query
+    - **query**: The search query
+    - **num_results**: Number of results to return
+    """
     try:
         document_list = load_embeddings(request.index_id)
         results = embeddings.search(request.query, request.num_results)