simplify the fast api example & upgrade query handling (#537)

Author: badmonster0

examples/fastapi_server_docker/.env

@@ -1 +1,5 @@
-COCOINDEX_DATABASE_URL=postgres://cocoindex:cocoindex@coco_db:5432/cocoindex
+# for docker
+COCOINDEX_DATABASE_URL=postgres://cocoindex:cocoindex@coco_db:5436/cocoindex
+
+# For local testing
+# COCOINDEX_DATABASE_URL=postgres://cocoindex:cocoindex@localhost/cocoindex

examples/fastapi_server_docker/README.md

@@ -1,10 +1,65 @@
-## Run cocoindex docker container with a simple query endpoint via fastapi
-In this example, we provide a simple docker container using docker compose to build pgvector17 along with a simple python fastapi script than runs a simple query endpoint. This example uses the code from the code embedding example.
+## Run docker container with a simple query endpoint via fastapi
 
-## How to run
-Edit the sample code directory to include the code you want to query over in
-```sample_code/``` 
+In this example, we will build index for text embedding from local markdown files, and provide a simple query endpoint via fastapi.
+We provide a simple docker container using docker compose to build pgvector17 along with a simple python fastapi script
 
-Edit the configuration code from the file ```src/cocoindex_funs.py``` line 23 to 25.
+We appreciate a star ⭐ at [CocoIndex Github](https://github.com/cocoindex-io/cocoindex) if this is helpful.
 
-Finally build the docker container via: ```docker compose up``` while inside the directory of the example.
+
+## Run locally without docker
+
+In the `.env` file, use local Postgres URL
+
+```
+# For local testing
+COCOINDEX_DATABASE_URL=postgres://cocoindex:cocoindex@localhost/cocoindex
+```
+
+- Install dependencies:
+
+    ```bash
+    pip install -e .
+    ```
+
+- Setup:
+
+    ```bash
+    cocoindex setup main.py
+    ```
+
+- Update index:
+
+    ```bash
+    cocoindex update main.py
+    ```
+
+- Run:
+
+    ```bash
+    uvicorn main:fastapi_app --reload --host 0.0.0.0 --port 8000
+    ```
+ 
+ ## Query the endpoint
+
+    ```bash
+    curl "http://localhost:8000/search?q=model&limit=3"
+    ```
+
+
+## Run Docker
+
+In the `.env` file, use Docker Postgres URL
+
+```
+COCOINDEX_DATABASE_URL=postgres://cocoindex:cocoindex@coco_db:5436/cocoindex
+```
+
+Build the docker container via: 
+```bash
+docker compose up --build
+```
+
+Test the endpoint:
+```bash
+curl "http://0.0.0.0:8080/search?q=model&limit=3"
+```

examples/fastapi_server_docker/compose.yaml

@@ -6,8 +6,10 @@ services:
       POSTGRES_USER: cocoindex
       POSTGRES_PASSWORD: cocoindex
       POSTGRES_DB: cocoindex
+      POSTGRES_PORT: 5436
     ports:
-      - "5432:5432"
+      - "5436:5436"
+    command: postgres -p 5436
 
   coco_api:
     build:

examples/fastapi_server_docker/dockerfile

@@ -2,6 +2,12 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
+# Install PostgreSQL client libraries
+RUN apt-get update && apt-get install -y \
+    libpq-dev \
+    gcc \
+    && rm -rf /var/lib/apt/lists/*
+
 COPY requirements.txt .
 
 RUN pip install -r requirements.txt

examples/fastapi_server_docker/files/1810.04805v2.md

@@ -1,25 +1,82 @@
 import cocoindex
 import uvicorn
-
-from fastapi import FastAPI
 from dotenv import load_dotenv
+from fastapi import FastAPI, Query
+from psycopg_pool import ConnectionPool
+import os
+
+@cocoindex.transform_flow()
+def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
+    """
+    Embed the text using a SentenceTransformer model.
+    This is a shared logic between indexing and querying.
+    """
+    return text.transform(
+        cocoindex.functions.SentenceTransformerEmbed(
+            model="sentence-transformers/all-MiniLM-L6-v2"))
+
+@cocoindex.flow_def(name="MarkdownEmbeddingFastApiExample")
+def markdown_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
+    """
+    Define an example flow that embeds markdown files into a vector database.
+    """
+    data_scope["documents"] = flow_builder.add_source(
+        cocoindex.sources.LocalFile(path="files"))
+    doc_embeddings = data_scope.add_collector()
+
+    with data_scope["documents"].row() as doc:
+        doc["chunks"] = doc["content"].transform(
+            cocoindex.functions.SplitRecursively(),
+            language="markdown", chunk_size=2000, chunk_overlap=500)
+        
+        with doc["chunks"].row() as chunk:
+            chunk["embedding"] = text_to_embedding(chunk["text"])
+            doc_embeddings.collect(
+                filename=doc["filename"], 
+                location=chunk["location"],
+                text=chunk["text"], 
+                embedding=chunk["embedding"]
+            )
+
+    doc_embeddings.export(
+        "doc_embeddings",
+        cocoindex.storages.Postgres(),
+        primary_key_fields=["filename", "location"],
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 
-from src.cocoindex_funs import code_embedding_flow, code_to_embedding
+def search(pool: ConnectionPool, query: str, top_k: int = 5):
+    # Get the table name, for the export target in the text_embedding_flow above.
+    table_name = cocoindex.utils.get_target_storage_default_name(markdown_embedding_flow, "doc_embeddings")
+    # Evaluate the transform flow defined above with the input query, to get the embedding.
+    query_vector = text_to_embedding.eval(query)
+    # Run the query and get the results.
+    with pool.connection() as conn:
+        with conn.cursor() as cur:
+            cur.execute(f"""
+                SELECT filename, text, embedding <=> %s::vector AS distance
+                FROM {table_name} ORDER BY distance LIMIT %s
+            """, (query_vector, top_k))
+            return [
+                {"filename": row[0], "text": row[1], "score": 1.0 - row[2]}
+                for row in cur.fetchall()
+            ]
 
 fastapi_app = FastAPI()
-    
-query_handler = cocoindex.query.SimpleSemanticsQueryHandler(
-    name="SemanticsSearch",
-    flow=code_embedding_flow,
-    target_name="code_embeddings",
-    query_transform_flow=code_to_embedding,
-    default_similarity_metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY
-)
-
-@fastapi_app.get("/query")
-def query_endpoint(string: str):
-    results, _ = query_handler.search(string, 10)
-    return results
+
+@fastapi_app.on_event("startup")
+def startup_event():
+    load_dotenv()
+    cocoindex.init()
+    # Initialize database connection pool
+    fastapi_app.state.pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL"))
+
+@fastapi_app.get("/search")
+def search_endpoint(q: str = Query(..., description="Search query"), limit: int = Query(5, description="Number of results")):
+    results = search(fastapi_app.state.pool, q, limit)
+    return {"results": results}
 
 if __name__ == "__main__":
     load_dotenv()

examples/fastapi_server_docker/main.py

@@ -2,4 +2,6 @@ cocoindex>=0.1.42
 python-dotenv>=1.0.1
 fastapi==0.115.12
 fastapi-cli==0.0.7
-uvicorn==0.34.2
+uvicorn==0.34.2
+psycopg==3.2.6
+psycopg_pool==3.2.6