Add more API routes

pamelafox · pamelafox · commit 37901c54bc11 · 2024-05-15T14:25:28.000Z
diff --git a/.devcontainer/docker-compose.yaml b/.devcontainer/docker-compose.yaml
@@ -24,7 +24,7 @@ services:
       - postgres-data:/var/lib/postgresql/data
     environment:
       POSTGRES_DB: postgres
-      POSTGRES_USER: postgres
+      POSTGRES_USER: admin
       POSTGRES_PASSWORD: postgres
 
     # Add "forwardPorts": ["5432"] to **devcontainer.json** to forward PostgreSQL locally.
diff --git a/.env.sample b/.env.sample
@@ -1,6 +1,6 @@
 # Use these values to connect to the local database from within the devcontainer
 POSTGRES_HOST=localhost
-POSTGRES_USERNAME=postgres
+POSTGRES_USERNAME=admin
 POSTGRES_PASSWORD=postgres
 POSTGRES_DATABASE=postgres
 POSTGRES_SSL=disable
diff --git a/.github/workflows/app-tests.yaml b/.github/workflows/app-tests.yaml
@@ -19,6 +19,7 @@ jobs:
       postgres:
         image: pgvector/pgvector:pg16
         env:
+          POSTGRES_USER: admin
           POSTGRES_PASSWORD: postgres
         ports:
         - 5432:5432
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -7,7 +7,7 @@
             "driver": "PostgreSQL",
             "name": "local",
             "database": "postgres",
-            "username": "postgres",
+            "username": "admin",
             "password": "postgres"
         },
         {
diff --git a/src/fastapi_app/api_routes.py b/src/fastapi_app/api_routes.py
@@ -1,14 +1,41 @@
 import fastapi
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker
 
-from .api_models import ChatRequest
-from .globals import global_storage
-from .postgres_searcher import PostgresSearcher
-from .rag_advanced import AdvancedRAGChat
-from .rag_simple import SimpleRAGChat
+from fastapi_app.api_models import ChatRequest
+from fastapi_app.globals import global_storage
+from fastapi_app.postgres_models import Item
+from fastapi_app.postgres_searcher import PostgresSearcher
+from fastapi_app.rag_advanced import AdvancedRAGChat
+from fastapi_app.rag_simple import SimpleRAGChat
 
 router = fastapi.APIRouter()
 
 
+@router.get("/items/{id}")
+async def item_handler(id: int):
+    """A simple API to get an item by ID."""
+    async_session_maker = async_sessionmaker(global_storage.engine, expire_on_commit=False)
+    async with async_session_maker() as session:
+        item = (await session.scalars(select(Item).where(Item.id == id))).first()
+        return item.to_dict()
+
+
+@router.get("/similar")
+async def similar_handler(id: int, n: int = 5):
+    """A similarity API to find items similar to items with given ID."""
+    async_session_maker = async_sessionmaker(global_storage.engine, expire_on_commit=False)
+    async with async_session_maker() as session:
+        item = (await session.scalars(select(Item).where(Item.id == id))).first()
+        closest = await session.execute(
+            select(Item, Item.embedding.l2_distance(item.embedding))
+            .filter(Item.id != id)
+            .order_by(Item.embedding.l2_distance(item.embedding))
+            .limit(n)
+        )
+        return [item.to_dict() | {"distance": round(distance, 2)} for item, distance in closest]
+
+
 @router.post("/chat")
 async def chat_handler(chat_request: ChatRequest):
     messages = [message.model_dump() for message in chat_request.messages]

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@`
`7`	`7`	`"driver": "PostgreSQL",`
`8`	`8`	`"name": "local",`
`9`	`9`	`"database": "postgres",`
`10`		`- "username": "postgres",`
	`10`	`+ "username": "admin",`
`11`	`11`	`"password": "postgres"`
`12`	`12`	`},`
`13`	`13`	`{`