46 lines
1.3 KiB
Python
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 |