Add Haystack

.devcontainer/devcontainer.json

@@ -83,5 +83,8 @@
     "forwardPorts": [
         8000, 50051, 8765
     ],
+    "mounts": [
+        "source=d:/Work/ya.practicum/llm,target=/var/run/llm,type=bind"
+    ],
     "postCreateCommand": "find . -name \"requirements.txt\" -type f -exec pip install -r '{}' ';'"
 }

assistant/README.md

@@ -0,0 +1,11 @@
+# RAG (Retrieval-Augmented Generation)
+
+## Options
+[LangChain](https://python.langchain.com/docs/integrations/vectorstores/elasticsearch/): A library for building RAG pipelines.
+[Haystack](https://haystack.deepset.ai/integrations/elasticsearch-document-store): Specifically designed for search and question answering systems
+
+## LLM model
+Clone the repo from the [link](https://huggingface.co/deepset/roberta-base-squad2)
+
+## Warning
+None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.

assistant/requirements.txt

@@ -0,0 +1,9 @@
+fastapi==0.110.1
+pydantic==2.6.4
+pydantic-settings~=2.2.1
+elasticsearch-haystack~=1.0.1
+uvicorn==0.29.0
+uvloop==0.19.0 ; sys_platform != "win32" and implementation_name == "cpython"
+tensorflow~=2.18.0
+sentence-transformers~=3.3.1
+tf-keras~=2.18.0

assistant/src/core/settings.py

@@ -0,0 +1,14 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class ElasticsearchSettings(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="ELASTIC_")
+    url: str = ""
+
+    @property
+    def host(self) -> str:
+        return self.url.rsplit(":", maxsplit=1)[0]
+
+    @property
+    def port(self) -> int:
+        return int(self.url.rsplit(":", maxsplit=1)[1])

assistant/src/main.py

@@ -0,0 +1,54 @@
+import os
+
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from src.pipeline import pipeline
+
+app = FastAPI(
+    title="Assistant",
+    description="Description",
+    version="1.0.0",
+    docs_url="/api/openapi",
+    openapi_url="/api/openapi.json",
+)
+try:
+    file_size = os.path.getsize("/var/run/llm/roberta-base-squad2/tf_model.h5")
+    print(file_size)
+except Exception as e:
+    print(e)
+
+
+# Define a request schema
+class QueryRequest(BaseModel):
+    query: str
+
+
+@app.post("/search")
+async def search(query_request: QueryRequest):
+    query = {"text_embedder": {"text": query_request.query}}
+    try:
+        # Perform the query using Haystack
+        result = pipeline.run(query)
+        return {
+            "query": query,
+            "answers": [
+                {
+                    "answer": answer.answer,
+                    "score": answer.score,
+                    "context": answer.context,
+                }
+                for answer in result["answers"]
+            ],
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    uvicorn.run(
+        "main:app",
+        host="0.0.0.0",
+        port=8000,
+        reload=True,
+    )

assistant/src/pipeline.py

@@ -0,0 +1,30 @@
+import os
+
+from haystack import Pipeline
+from haystack.components.embedders import SentenceTransformersTextEmbedder
+from haystack_integrations.components.retrievers.elasticsearch import (
+    ElasticsearchEmbeddingRetriever,
+)
+from haystack_integrations.document_stores.elasticsearch import (
+    ElasticsearchDocumentStore,
+)
+from src.core.settings import ElasticsearchSettings
+
+# Initialize Elasticsearch Document Store
+settings = ElasticsearchSettings()
+index = os.environ["ELASTIC_INDEX_NAME_MOVIES"]
+document_store = ElasticsearchDocumentStore(hosts=settings.url, index=index)
+
+# Initialize Retriever
+retriever = ElasticsearchEmbeddingRetriever(document_store=document_store)
+
+# Optionally add a Reader for better answers
+# if nothing specified, model will be downloaded on startup
+model = os.getenv("ASSISTANT_MODEL", "deepset/roberta-base-squad2")
+text_embedder = SentenceTransformersTextEmbedder(model=model)
+
+# Combine Retriever and Reader into a Pipeline
+pipeline = Pipeline()
+pipeline.add_component("text_embedder", text_embedder)
+pipeline.add_component("retriever", retriever)
+pipeline.connect("text_embedder.embedding", "retriever.query_embedding")

infra/.env.example

@@ -92,4 +92,8 @@ [email protected]
 MAILHOG_ENABLED=True
 MAILHOG_HOST=mailhog
 MAILHOG_PORT=1025
-[email protected]
+[email protected]
+
+# Assistant
+# it can be either name or path
+ASSISTANT_MODEL=/path/to/roberta-base-squad2

infra/docker-compose.yaml

@@ -69,7 +69,7 @@ services:
   elasticsearch-dev:
     <<: *elasticsearch-base
     ports: ["9200:9200"]
-    profiles: ["etl-dev", "api-dev"]
+    profiles: ["etl-dev", "api-dev", "assist-dev"]
 
   elasticsearch-clean:
     <<: *elasticsearch-base