Update app.py

app.py

@@ -1,54 +1,103 @@
 import os
+import asyncio
+import json
 import gradio as gr
-from huggingface_hub import InferenceClient
-from tinyagents import Agent
+from contextlib import AsyncExitStack
+from typing import Optional
 
-# Set your Hugging Face API token
-HF_TOKEN = os.getenv("HF_TOKEN")
+from huggingface_hub import AsyncInferenceClient
+from mcp import ClientSession
+from mcp.http import MCPServerHTTP  # Use HTTP MCP client for remote server
+
+SYSTEM_PROMPT = """Your name is Jobcy. You are an AI assistant designed to help users find remote jobs by searching through job listings from various sources, including the Jobicy API and other platforms.
+You will list the job listings in a structured format, including the job title, company, location, and a google search link."""
 
-# Define the MCP server URL
 MCP_URL = "https://oppaai-job-search-mcp-server.hf.space/gradio_api/mcp/sse"
 
-# Initialize the MCP client
-mcp_client = InferenceClient(
-    url=MCP_URL,
-    token=HF_TOKEN
-)
-
-# Fetch the tools exposed by the MCP server
-tools = mcp_client.get_tools()
-
-# Define the model to be used
-model = InferenceClient(
-    model="Qwen/Qwen3-30B-A3B",
-    token=HF_TOKEN
-)
-
-# Create the agent
-agent = Agent(
-    model=model,
-    tools=tools
-)
-
-# Define the function to handle user input and generate responses
-def chat_fn(user_message, history):
+class TinyToolCallingAgent:
+    def __init__(self):
+        self.mcp_client_session: Optional[ClientSession] = None
+        self.exit_stack = AsyncExitStack()
+        self.model = AsyncInferenceClient(model="Qwen/Qwen3-8B", provider="hf-inference", token=os.getenv("HF_TOKEN"))
+        self.system_prompt = SYSTEM_PROMPT
+        self.tools = []
+
+    async def connect_to_server(self):
+        # Connect to MCP server over HTTP SSE URL
+        server = MCPServerHTTP(url=MCP_URL)
+        self.mcp_client_session = await server.connect()
+        response = await self.mcp_client_session.list_tools()
+        self.tools = response.tools
+        print("Connected to MCP server with tools:", [tool.name for tool in self.tools])
+
+    async def _call_model(self, messages, tools=None):
+        return await self.model.chat_completion(messages, max_tokens=1000, tools=tools)
+
+    async def process_query(self, query: str) -> str:
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": query},
+        ]
+
+        tools = [
+            {
+                "type": "function",
+                "function": {"name": tool.name, "description": tool.description, "parameters": tool.inputSchema},
+            }
+            for tool in self.tools
+        ]
+
+        response = await self._call_model(messages, tools=tools)
+        final_text = []
+        message = response.choices[0].message
+
+        if message.tool_calls:
+            if message.content:
+                messages.append({"role": "assistant", "content": message.content})
+            for tool_call in message.tool_calls:
+                tool_name = tool_call.function.name
+                tool_args = json.loads(tool_call.function.arguments)
+                result = await self.mcp_client_session.call_tool(tool_name, tool_args)
+                final_text.append(f"[Calling tool {tool_name} with args {tool_args}]")
+                messages.append({"role": "user", "content": result.content[0].text})
+            response = await self._call_model(messages)
+            final_text.append(response.choices[0].message.content)
+        elif message.content:
+            final_text.append(message.content)
+
+        return "\n".join(final_text)
+
+    async def cleanup(self):
+        if self.mcp_client_session:
+            await self.mcp_client_session.aclose()
+        await self.exit_stack.aclose()
+
+agent = TinyToolCallingAgent()
+
+async def gradio_chat(user_input, history):
     history = history or []
-    history.append(("user", user_message))
-    response = agent.run(user_message)
-    history.append(("agent", response))
+    try:
+        response = await agent.process_query(user_input)
+    except Exception as e:
+        response = f"Error: {str(e)}"
+    history.append(("User", user_input))
+    history.append(("Jobcy", response))
     return history, history
 
-# Set up the Gradio interface
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    user_input = gr.Textbox(placeholder="Ask Jobcy about remote jobs or anything else...")
-    state = gr.State([])
+async def main():
+    await agent.connect_to_server()
+
+    with gr.Blocks() as demo:
+        chatbot = gr.Chatbot()
+        user_input = gr.Textbox(placeholder="Ask Jobcy about remote jobs...", label="Your Message")
+        state = gr.State([])
+
+        user_input.submit(gradio_chat, inputs=[user_input, state], outputs=[chatbot, state])
 
-    user_input.submit(chat_fn, inputs=[user_input, state], outputs=[chatbot, state])
+        demo.title = "Jobcy Remote Job Search Assistant"
+        demo.launch()
 
-demo.title = "Jobcy Remote Job Search Assistant (tinyagents)"
-demo.launch()
+    await agent.cleanup()
 
-# Disconnect the MCP client when done
-import atexit
-atexit.register(mcp_client.disconnect)
+if __name__ == "__main__":
+    asyncio.run(main())