Update app.py

app.py

@@ -1,103 +1,56 @@
-import os
 import asyncio
-import json
+import re
+import os
 import gradio as gr
-from contextlib import AsyncExitStack
-from typing import Optional
-
-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."""
-
-MCP_URL = "https://oppaai-job-search-mcp-server.hf.space/gradio_api/mcp/sse"
-
-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):
+from pydantic_ai import Agent
+from pydantic_ai.mcp import MCPServerHTTP
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.openai import OpenAIProvider
+
+# MCP Server SSE URL
+SSE_URL = "https://oppaai-job-search-mcp-server.hf.space/gradio_api/mcp/sse"
+server = MCPServerHTTP(url=SSE_URL)
+
+HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure your HF_TOKEN is set in environment variables
+
+# Use HuggingFace hosted Qwen3-30B-A3B model
+qwen3_model = OpenAIModel(
+    model_name="Qwen/Qwen3-30B-A3B",
+    provider=OpenAIProvider(
+        base_url="https://api-inference.huggingface.co/models/Qwen/Qwen3-30B-A3B",
+        token=HF_TOKEN,
+    ),
+)
+
+# Create Agent with MCP Server
+agent = Agent(
+    model=qwen3_model,
+    mcp_servers=[server],
+    instructions="""
+    Your name is Jobcy. You are an AI assistant designed to help users to 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 the google search link.
+    """
+)
+
+async def chat_with_agent(user_input, history):
     history = history or []
-    try:
-        response = await agent.process_query(user_input)
-    except Exception as e:
-        response = f"Error: {str(e)}"
+    result = await agent.run(user_input)
+    cleaned_output = re.sub(r"<think>.*?</think>", "", result.output, flags=re.DOTALL).strip()
     history.append(("User", user_input))
-    history.append(("Jobcy", response))
+    history.append(("Jobcy", cleaned_output))
     return history, history
 
 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])
+    async with agent.run_mcp_servers():
+        with gr.Blocks() as demo:
+            chatbot = gr.Chatbot()
+            user_input = gr.Textbox(placeholder="Ask Jobcy about remote jobs or anything else...", label="Your Message")
+            state = gr.State([])
 
-        demo.title = "Jobcy Remote Job Search Assistant"
-        demo.launch()
+            user_input.submit(chat_with_agent, inputs=[user_input, state], outputs=[chatbot, state])
 
-    await agent.cleanup()
+            demo.title = "Jobcy Remote Job Search Assistant"
+            demo.launch()
 
 if __name__ == "__main__":
     asyncio.run(main())