app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os
from dotenv import load_dotenv
from llama_index.core.indices.vector_store.base import VectorStoreIndex
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.embeddings.fastembed import FastEmbedEmbedding
from llama_index.core import Settings
import qdrant_client
from llama_index.llms.gemini import Gemini
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.readers.web import FireCrawlWebReader
from fastapi.middleware.cors import CORSMiddleware
from llama_index.core.agent import ReActAgent
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.core.llms import ChatMessage
from composio_llamaindex import ComposioToolSet, App, Action 
from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.llms import ChatMessage
from llama_index.llms.openai import OpenAI
from pathlib import Path

class AgentInput(BaseModel):
    sheet_id: str
    api_key: str
    query: str




import nltk
import ssl

try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    pass
else:
    ssl._create_default_https_context = _create_unverified_https_context

# Set the NLTK data path explicitly
nltk.data.path.append(os.environ["NLTK_DATA"])

# Download specific NLTK datasets
nltk.download('punkt', download_dir=os.environ["NLTK_DATA"])

load_dotenv()

from fastapi.middleware.httpsredirect import HTTPSRedirectMiddleware

app = FastAPI()
#app.add_middleware(HTTPSRedirectMiddleware)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust this to your needs
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize session state equivalent
state = {
    'setup_complete': False,
    'documents': None,
    'chat_history': [],
    'index': None,
    'url': "",
    'collection_name': "",
}

os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
os.environ["COMPOSIO_API_KEY"] = os.getenv("COMPOSIO_API_KEY")

# Setup functions
def embed_setup():
    Settings.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-en-v1.5")
    Settings.llm = Gemini(temperature=0.1, model_name="models/gemini-1.5-pro")

@app.post("/sheets_query/")
async def process_query(input_data: AgentInput):
    # Load environment variables
    load_dotenv()

    # Set up OpenAI LLM
    #os.environ["OPENAI_API_KEY"] = 
    llm = OpenAI(model='gpt-4o', api_key=input_data.api_key)

    # Set up ComposioToolSet
    composio_toolset = ComposioToolSet(api_key=os.getenv("COMPOSIO_API_KEY"), output_dir=Path("./plots/"))
    tools = composio_toolset.get_tools(apps=[App.GOOGLESHEETS, App.CODEINTERPRETER])

    # Define prefix messages
    prefix_messages = [
        ChatMessage(
            role="system",
            content=(
                "You are an AI assistant who is an expert at Google Sheets. "
                "Use Google Sheets Tool and perform the necessary operations based on query, use code interpreter to plot graphs"
                "return the plotted graph images in markdown format."
                "create a sandbox and use the codeinterpreter."
            )
        )
    ]

    # Create agent
    agent = FunctionCallingAgentWorker(
        tools=tools,
        llm=llm,
        prefix_messages=prefix_messages,
        max_function_calls=10,
        allow_parallel_tool_calls=False,
        verbose=True
    ).as_agent()

    try:
        # Process the query
        response = agent.chat(f"This is the Google Sheet ID: {input_data.sheet_id}. {input_data.query}")
        return {"response": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

        
def qdrant_setup():
    client = qdrant_client.QdrantClient(
        os.getenv("QDRANT_URL"),
        api_key=os.getenv("QDRANT_API_KEY"),
    )
    return client

def llm_setup():
    llm = Gemini(api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.3, model_name="models/gemini-pro")
    return llm

def ingest_documents(url):
    firecrawl_reader = FireCrawlWebReader(
        api_key=os.getenv("FIRECRAWL_API_KEY"),
        mode="scrape",
    )
    documents = firecrawl_reader.load_data(url=url)
    return documents

class SetupRequest(BaseModel):
    url: str = ""
    collection_name: str

class QueryRequest(BaseModel):
    query: str

@app.post("/setup/")
async def setup(request: SetupRequest):
    
    state['url'] = request.url
    state['collection_name'] = request.collection_name

    embed_setup()
    client = qdrant_setup()
    llm = llm_setup()
    vector_store = QdrantVectorStore(client=client, collection_name=state['collection_name'])

    if state['url']:
        state['documents'] = ingest_documents(state['url'])
        state['index'] = VectorStoreIndex.from_documents(state['documents'], vector_store=vector_store)
        state['setup_complete'] = True
        return {"message": f"Documents ingested from {state['url']} and query engine setup completed successfully!"}
    else:
        state['index'] = VectorStoreIndex.from_vector_store(vector_store=vector_store)
        state['setup_complete'] = True
        return {"message": f"Query engine setup completed successfully using existing collection: {state['collection_name']}"}

@app.post("/query/")
async def query(request: QueryRequest):
    if not state['setup_complete']:
        raise HTTPException(status_code=400, detail="Please complete the setup first")
    
    memory = ChatMemoryBuffer.from_defaults(token_limit=4000)
    chat_engine = state['index'].as_chat_engine(
        chat_mode="context",
        memory=memory,
        system_prompt=(
            """You are an AI assistant for developers, specializing in technical documentation. You are also great at casual conversations.
                Your task is to provide accurate, concise, and helpful responses based on the given documentation context.
                Context information is below:
                {context_str}
                Always answer based on the information in the context and general knowledge and be precise
                Given this context, please respond to the following user query:
                {query_str}
                Your response should:
                Directly address the query using information from the context
                Include relevant code examples or direct quotes if applicable
                Mention specific sections or pages of the documentation
                Highlight any best practices or potential pitfalls related to the query
                After your response, suggest 3 follow-up questions based on the context that the user might find helpful for deeper understanding.
                
                ALWAYS SUGGEST FOLLOW UP QUESTIONS
                Your response:"""
        ),
    )
    
    response = chat_engine.chat(request.query)
    state['chat_history'].append(("User", request.query))
    state['chat_history'].append(("Assistant", str(response.response)))

    return {"response": response.response}


@app.post("/agent-search/")
async def agent_search(request: QueryRequest):
    code_tools=DuckDuckGoSearchToolSpec()
    tool = code_tools.to_tool_list()
    llm = Gemini(api_key=os.environ["GOOGLE_API_KEY"], model="models/gemini-1.5-pro")
    prefix_messages = [
        ChatMessage(
            role="system",
            content=(
                "You are now a integration agent, and what  ever you are requested, you will try to execute utilizing your toools."
            ),
        )
    ]
    agent = ReActAgent.from_tools(tool, llm=llm, verbose=True)
    response = agent.chat(request.query)
    state['chat_history'].append(("User", request.query))
    state['chat_history'].append(("Assistant", str(response.response)))

    return {"response": response.response}


@app.get("/chat-history/")
async def get_chat_history():
    return {"chat_history": state['chat_history']}



@app.post("/clear-chat/")
async def clear_chat():
    state['chat_history'] = []
    return {"message": "Chat history cleared!"}