diff --git a/entrypoint.py b/entrypoint.py index 3b1b0c5..5ae7ea6 100644 --- a/entrypoint.py +++ b/entrypoint.py @@ -1,85 +1,46 @@ -# main.py -from qdrant_client.http import models as qmodels -import uuid -from fastapi import FastAPI -import requests, os, redis -from qdrant_client import QdrantClient -from langchain_community.vectorstores import Qdrant -from langchain_community.embeddings import HuggingFaceEmbeddings +import os +import langchain +from langchain.agents import tool +from langchain.tools import searxng, qdrant, llma_server +# Environment variables +SEARXNG_URL = os.environ['SEARXNG_URL'] +QDRANT_URL = os.environ['QDRANT_URL'] +LLAMA_SERVER_URL = os.environ['LLAMA_SERVER_URL'] + +# Define the LangChain agent +class SearchAgent(langchain.Agent): + def __init__(self): + super().__init__() + self.searxng_tool = searxng.SearXNGTool(SEARXNG_URL=SEARXNG_URL) + self.qdrant_tool = qdrant.QDrantTool(QDRANT_URL=QDRANT_URL) + self.llma_server_tool = llma_server.LLaMAServerTool(LLAMA_SERVER_URL=LLAMA_SERVER_URL) + + async def search(self, query): + # Send query to SearXNG API + results = await self.searxng_tool.search(query) + # Embed results with QDrant + await self.qdrant_tool.embed(results) + return results + + async def query(self, query): + # Use LLaMA-Server to process query + response = await self.llma_server_tool.query(query) + return response + +# Create the LangChain agent +agent = SearchAgent() + +# Define the API endpoints +from fastapi import FastAPI app = FastAPI() -# Connectors -qdrant = QdrantClient(url=os.getenv("QDRANT_URL")) -r = redis.Redis.from_url(os.getenv("REDIS_URL")) -SEARXNG_URL = os.getenv("SEARXNG_URL") -LLAMA_SERVER_URL = os.getenv("LLAMA_SERVER_URL") -QDRANT_URL = os.getenv("QDRANT_URL") -REDIS_URL = os.getenv("REDIS_URL") - -embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") - -@app.get("/ask") -def ask(query: str): - # 1. Search via SearxNG - resp = requests.get(f"{SEARXNG_URL}/search", params={"q": query, "format": "json"}) - snippets = [r["title"] + " " + r["content"] for r in resp.json()["results"][:3]] - - context = " ".join(snippets) - - # 2. Embed + store in Qdrant - vectors = embeddings.embed_documents(snippets) - points = [ - qmodels.PointStruct( - id=str(uuid.uuid4()), # unique ID for each snippet - vector=vec, # the embedding vector - payload={"text": text} # optional metadata - ) - for text, vec in zip(snippets, vectors) - ] - qdrant.upsert(collection_name="docs", points=points) - - # 3. Call llama-server - llama_http = requests.post( - f"{LLAMA_SERVER_URL}/completion", - json={ - "prompt": f"Context:\n{context}\n\nQuestion: {query}\nAnswer:", - "n_predict": 128, - "temperature": 0.7, - "stop": [""] - } - ) - llm_resp = llama_http.json() - print("Llama raw response:", llm_resp) - - return {"answer": llm_resp.get("content", "")} -@app.get("/health") -def health(): - results = {} - - # Check Redis - try: - r = redis.Redis.from_url(REDIS_URL) - r.ping() - results["redis"] = "ok" - except Exception as e: - results["redis"] = f"error: {e}" - - # Check Qdrant - try: - q = requests.get(f"{QDRANT_URL}/readyz") - results["qdrant"] = f"ok ({q.status_code})" - except Exception as e: - results["qdrant"] = f"error: {e}" - - # Check SearxNG - try: - s = requests.get(f"{SEARXNG_URL}/search", - params={"q": "ping", "format": "json"}, - headers={"X-Forwarded-For": "127.0.0.1"}) - results["searxng"] = f"ok ({s.status_code})" - except Exception as e: - results["searxng"] = f"error: {e}" - +@app.post("/search") +async def search(query: str): + results = await agent.search(query) return results +@app.post("/query") +async def query(query: str): + response = await agent.query(query) + return response \ No newline at end of file