Update entrypoint.py
This commit is contained in:
parent
666680babb
commit
4fec3f6596
121
entrypoint.py
121
entrypoint.py
|
|
@ -1,85 +1,46 @@
|
||||||
# main.py
|
import os
|
||||||
from qdrant_client.http import models as qmodels
|
import langchain
|
||||||
import uuid
|
from langchain.agents import tool
|
||||||
from fastapi import FastAPI
|
from langchain.tools import searxng, qdrant, llma_server
|
||||||
import requests, os, redis
|
|
||||||
from qdrant_client import QdrantClient
|
|
||||||
from langchain_community.vectorstores import Qdrant
|
|
||||||
from langchain_community.embeddings import HuggingFaceEmbeddings
|
|
||||||
|
|
||||||
|
# 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()
|
app = FastAPI()
|
||||||
|
|
||||||
# Connectors
|
@app.post("/search")
|
||||||
qdrant = QdrantClient(url=os.getenv("QDRANT_URL"))
|
async def search(query: str):
|
||||||
r = redis.Redis.from_url(os.getenv("REDIS_URL"))
|
results = await agent.search(query)
|
||||||
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": ["</s>"]
|
|
||||||
}
|
|
||||||
)
|
|
||||||
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}"
|
|
||||||
|
|
||||||
return results
|
return results
|
||||||
|
|
||||||
|
@app.post("/query")
|
||||||
|
async def query(query: str):
|
||||||
|
response = await agent.query(query)
|
||||||
|
return response
|
||||||
Loading…
Reference in New Issue