Update app.py

app.py

@@ -14,6 +14,16 @@ 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
 
 
 
@@ -66,6 +76,48 @@ 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"] = input_data.api_key
+    llm = OpenAI(model='gpt-4o')
+
+    # 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])
+
+    # 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"
+            )
+        )
+    ]
+
+    # 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"),
@@ -174,6 +226,8 @@ async def agent_search(request: QueryRequest):
 async def get_chat_history():
     return {"chat_history": state['chat_history']}
 
+
+
 @app.post("/clear-chat/")
 async def clear_chat():
     state['chat_history'] = []