init vector db api

services/backend/vector_db_api/Dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.8-slim-buster
+
+# Create a separate directory for dependencies
+WORKDIR /app/requirements
+
+# Copy requirements.txt into the container
+COPY requirements.txt .
+
+# Install the required packages
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Set the working directory
+WORKDIR /app
+
+# Copy the FastAPI app script into the container
+COPY milvus ./milvus
+
+
+# Expose the default port for FastAPI
+EXPOSE 8010
+
+ENV HOST=host.docker.internal
+
+# Run the FastAPI app using uvicorn
+CMD ["uvicorn", "milvus.main:app", "--host", "0.0.0.0", "--port", "8010"]

services/backend/vector_db_api/README.md

@@ -0,0 +1,11 @@
+# Vector Search Service
+
+This service provides a FastAPI application that interfaces with a vector database for text embedding and similarity search operations. It is designed to be flexible and allows for the potential integration of various vector databases in the future.
+
+The service provides the following endpoints:
+
+- `POST /add_texts`: Add texts to a specified Milvus collection
+- `DELETE /delete_docs`: Delete documents from a specified Milvus collection based on an expression
+- `POST /similarity_match`: Perform a similarity search and return the top-k similar documents
+- `GET /ping`: Health check endpoint
+

services/backend/vector_db_api/milvus/CustomEmbeddings.py

@@ -0,0 +1,55 @@
+# Standard Library
+import json
+import os
+from typing import List
+
+# Third Party
+import requests
+from langchain.embeddings.base import Embeddings
+
+HOST = os.getenv("EMBEDDING_HOST", "localhost")
+EMBEDDING_URL = f"http://{HOST}:8020/embed_multiple"
+
+
+def embed_text(texts: List[str]) -> List[List[float]]:
+    """Embed a list of texts using a remote service.
+
+    Args:
+        texts (List[str]): List of texts to be embedded.
+
+    Returns:
+        List[List[float]]: List of embedded texts.
+    """
+    headers = {"accept": "application/json", "Content-Type": "application/json"}
+    data = json.dumps(texts)
+
+    response = requests.post(EMBEDDING_URL, headers=headers, data=data)
+    response = response.json()
+
+    return response["embeddings"]
+
+
+class CustomEmbeddings(Embeddings):
+    """Custom embeddings class that uses a remote service for embedding."""
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed a list of documents.
+
+        Args:
+            texts (List[str]): List of documents to be embedded.
+
+        Returns:
+            List[List[float]]: List of embedded documents.
+        """
+        return embed_text(texts)
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed a single query.
+
+        Args:
+            text (str): Query to be embedded.
+
+        Returns:
+            List[float]: Embedded query.
+        """
+        return embed_text([text])[0]

services/backend/vector_db_api/milvus/main.py

@@ -0,0 +1,119 @@
+# Standard library imports
+# Standard Library
+import os
+from typing import List, Optional
+
+# Third Party
+# Third party imports
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+# Local application imports
+from .CustomEmbeddings import CustomEmbeddings
+from .query_milvus import Milvus
+
+MILVUS_HOST = os.getenv("MILVUS_HOST", "localhost")
+
+model = {}
+top_re_rank = 5
+top_k_match = 10
+app = FastAPI()
+
+
+class Query(BaseModel):
+    text: str
+    collection_name: str
+    topk: int = top_k_match
+    rerank: bool = False
+    topr: int = top_re_rank
+
+
+@app.on_event("startup")
+async def app_startup():
+    pass
+
+
+def drop_collection(collection_name: str):
+    load_collection(collection_name).drop_collection()
+
+
+def load_collection(collection_name: str) -> Milvus:
+    return Milvus(
+        embedding_function=CustomEmbeddings(),
+        collection_name=collection_name,
+        connection_args={
+            "host": MILVUS_HOST,
+            "port": "19530",
+            "user": "username",
+            "password": "password",
+        },
+        index_params={
+            "metric_type": "IP",
+            "index_type": "FLAT",
+            "params": {"nlist": 16384},
+        },
+        search_params={"metric_type": "IP", "params": {"nprobe": 32}},
+    )
+
+
+@app.post("/add_texts")
+async def add_texts_embeddings(
+    collection_name: str,
+    texts: List[str],
+    metadatas: Optional[List[dict]] = None,
+):
+    """_summary_
+
+    Args:
+        texts (List[str]): _description_
+        connlection_name (str): this should reflect user's random id + the assistant_id they created.
+    """
+    pks = add_texts(collection_name, texts, metadatas)
+
+
+def add_texts(
+    collection_name: str,
+    texts: List[str],
+    metadatas: Optional[List[dict]] = None,
+):
+    c = load_collection(collection_name)
+    pks = c.add_texts(texts=texts, metadatas=metadatas)
+    print(pks)
+    return pks
+
+
+@app.delete("/delete_docs")
+async def delete_docs_api(collection_name: str, expr: str):
+    delete_docs(collection_name, expr)
+
+
+def delete_docs(collection_name: str, expr: str):
+    c = load_collection(collection_name)
+    c.delete_entities(expr=expr)
+
+
+def get_top_k_biencoder_match_milvus(query: Query):
+    c = load_collection(query.collection_name)
+
+    docs = c.similarity_search(query.text, k=query.topk)
+    res = []
+    for i, d in enumerate(docs):
+        thisd = {"id": i, "metadata": d.metadata, "text": d.page_content}
+        res.append(thisd)
+    return res
+
+
+def get_similar_match(query, biencoder_match_method: str, rerank: bool = False):
+    query_biencoder_matches = get_top_k_biencoder_match_milvus(query)
+    return query_biencoder_matches[: query.topr]
+
+
+@app.post("/similarity_match")
+def text_similarity_match(query: Query):
+    res = get_similar_match(query, biencoder_match_method="milvus", rerank=query.rerank)
+    return {"response": res}
+
+
+@app.get("/ping")
+def ping():
+    return {"response": "Pong!"}