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