import asyncio
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from google import genai

client = genai.Client(api_key="<YOUR_GEMINI_API_KEY>")

async def run():
    async with streamablehttp_client(
        "https://mcp.trychroma.com/package-search/v1",
        headers={"x-chroma-token": "<YOUR_CHROMA_API_KEY>"},
    ) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()
            try:
                prompt = f"Explain how numpy implements its FFT. Use package search."
                response = await client.aio.models.generate_content(
                    model="gemini-2.5-flash",
                    contents=prompt,
                    config=genai.types.GenerateContentConfig(
                        temperature=0,
                        tools=[session],
                    ),
                )
                try:
                    if response.text:
                        print("--- Generated Text ---")
                        print(response.text)
                    else:
                        print("Model did not return text.")
                        print(f"Finish Reason: {response.candidates[0].finish_reason.name}")
                except ValueError:
                    print("Could not access response.text.")
            except Exception as e:
                print(f"An error occurred: {e}")

asyncio.run(run())