Desktop_AI_Agent/client.py at main · Anish-CodeDev/Desktop_AI_Agent · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import asyncio
from typing import Optional
from contextlib import AsyncExitStack
import re
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from google import genai
from dotenv import load_dotenv
load_dotenv()
gemini = genai.Client()


def get_result(query,tools):
    tools_with_desc = {
        tool.name:tool.description

    for tool in tools
    }
    response = gemini.models.generate_content(
            model="gemini-2.5-flash-lite",
            contents=f"""Given the {query}, decide the tool by the tool's name and description which are represented by triple backticks
            ```{tools_with_desc}```
            Choose the tool best suits the user's query, extract arguments,
            Return the best tool with the key 'tool'and the extracted arguments under the key 'args' in the form of json.
            Remove all the backticks like ``` at any cost.
            """,
    )
    return eval(str(re.sub('json','',response.text)))

class MCPClient:
    def __init__(self):
        self.session:Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()

    async def list_tools(self,loc):
        try:

            server_params = StdioServerParameters(
                command='python',
                args=[loc],
                env=None
            )
            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
            self.read,self.write = stdio_transport

            self.session = await self.exit_stack.enter_async_context(ClientSession(self.read,self.write))

            await self.session.initialize()

            tools_list = await self.session.list_tools()
            tools_list = tools_list.tools

        finally:
            await self.exit_stack.aclose()


        return {
             "name":[tool.name for tool in tools_list],
             "desc":[tool.description for tool in tools_list]
        }
    async def connect_with_server(self,query,loc):
        server_params = StdioServerParameters(
            command='python',
            args=[loc],
            env=None
        )
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.read,self.write = stdio_transport

        self.session = await self.exit_stack.enter_async_context(ClientSession(self.read,self.write))

        await self.session.initialize()


        response = await self.session.list_tools()
        tools = response.tools

        result = get_result(query,tools)

        return result

    async def call_tool(self,name,args,loc):
        server_params = StdioServerParameters(
            command='python',
            args=[loc],
            env=None
        )
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.read,self.write = stdio_transport

        self.session = await self.exit_stack.enter_async_context(ClientSession(self.read,self.write))

        await self.session.initialize()
        try:

            response = await self.session.call_tool(name,args)

        finally:
            await self.exit_stack.aclose()

        return response


client = MCPClient()

async def main(user_inp,loc):
    print(user_inp)
    try:
        #user_inp = input("User: ")
        #while user_inp !='exit':

            result = await client.connect_with_server(user_inp,loc=loc)
            tool_name = result['tool']
            args = result['args']
            print(tool_name)
            print(args)
            response = await client.session.call_tool(tool_name,args)

            #print("AI: ",response.content[0].text)
            #user_inp = input("User: ")
            return str(response)

    finally:
        await client.exit_stack.aclose()

#asyncio.run(main(''))