langchain/entrypoint.py

46 lines
1.3 KiB
Python

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()
@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