Couchbase MCP Server

import asyncio
import os
import shutil

from agents import Agent, Runner, gen_trace_id, trace, OpenAIChatCompletionsModel
from agents.mcp import MCPServer, MCPServerStdio
from openai import AsyncOpenAI
from dotenv import load_dotenv

load_dotenv()

api_key = os.environ["NEBIUS_API_KEY"]
base_url = "https://api.tokenfactory.nebius.com/v1"
client = AsyncOpenAI(base_url=base_url, api_key=api_key)

async def qna(agent):
    # Query 1: Database info
    message = "Tell me about the database that you are connected to."
    print(f"Running: {message}")
    result = await Runner.run(starting_agent=agent, input=message)
    print(result.final_output)

    # Query 2: Top hotels
    message = "List out the top 5 hotels by the highest aggregate rating?"
    print(f"\n\nRunning: {message}")
    result = await Runner.run(starting_agent=agent, input=message)
    print(result.final_output)

    # Query 3: Travel recommendation
    message = "Recommend me a flight and hotel from New York to San Francisco"
    print(f"\n\nRunning: {message}")
    result = await Runner.run(starting_agent=agent, input=message)
    print(result.final_output)

    # Query 4: UK sightseeing
    message = """I'm going to the UK for 1 week. Recommend some great spots 
    to visit for sightseeing. Also mention the respective prices of those 
    places for adults and kids."""
    print(f"\n\nRunning: {message}")
    result = await Runner.run(starting_agent=agent, input=message)
    print(result.final_output)

    # Query 5: Budget hotel
    message = "My budget is around 30 pounds a night. What will be the best hotel to stay in?"
    print(f"\n\nRunning: {message}")
    result = await Runner.run(starting_agent=agent, input=message)
    print(result.final_output)

async def run(mcp_server: MCPServer):
    agent = Agent(
        name="Assistant",
        instructions="""Couchbase organizes data with the following hierarchy:

        1. Cluster - Overall container of all data and services
        2. Bucket - Similar to a database (e.g., "travel-sample")
        3. Scope - Namespace within bucket (use "inventory" scope)
        4. Collection - Equivalent of a table (airport, airline, route, landmark, hotel)
        5. Document - JSON document with unique key

        Use the tools to read the database and answer questions.
        The data is inside `inventory` scope, so use only that scope.
        Any query you generate needs to have only the collection name in the FROM clause.
        Every field, collection, scope or bucket name inside the query should be 
        inside backticks (``).
        """,
        model=OpenAIChatCompletionsModel(
            model="meta-llama/Meta-Llama-3.1-8B-Instruct",
            openai_client=client
        ),
        mcp_servers=[mcp_server],
    )

    await qna(agent=agent)

async def main():
    # Configure MCP server
    async with MCPServerStdio(
        params={
            "command": "/path/to/uv",
            "args": [
                "--directory",
                "/path/to/mcp-server-couchbase",
                "run",
                "--env-file",
                ".env",
                "mcp_server.py"
            ]
        }
    ) as mcp_server:
        await run(mcp_server)

# Run the notebook
await main()

from agents.mcp import MCPServerStdio

async with MCPServerStdio(
    params={
        "command": "uv",  # Python environment manager
        "args": [
            "--directory",
            "/path/to/mcp-server-couchbase",  # Server location
            "run",
            "--env-file",
            ".env",  # Couchbase credentials
            "mcp_server.py"
        ]
    }
) as mcp_server:
    # Use mcp_server with agent

-- Correct: Only collection name in FROM
SELECT * FROM `hotel` WHERE `ratings`.`Overall` = 5.0

-- Incorrect: Including scope/bucket
SELECT * FROM `inventory`.`hotel`  -- Don't do this

-- All identifiers must use backticks
SELECT `name`, `city`, `ratings`.`Overall` 
FROM `hotel`
WHERE `country` = 'United States'

-- Top N results
SELECT `name`, `ratings`.`Overall` AS rating
FROM `hotel`
ORDER BY `ratings`.`Overall` DESC
LIMIT 5

-- Filtered results
SELECT `name`, `address`, `price`
FROM `hotel`
WHERE `country` = 'United Kingdom' 
  AND `price` <= 30

-- Join collections
SELECT h.`name`, a.`name` AS airport
FROM `hotel` h
JOIN `airport` a ON h.`city` = a.`city`
WHERE h.`country` = 'United States'

instructions = """
... (existing instructions) ...

Additional collections:
- `custom_collection`: Your custom data
- Use same query rules for custom collections
"""

instructions = """
Available scopes:
- inventory: Hotels, airlines, routes
- tenant_agent_00: User bookings

Always specify which scope to query.
"""

SELECT `country`, COUNT(*) AS hotel_count,
       AVG(`ratings`.`Overall`) AS avg_rating
FROM `hotel`
GROUP BY `country`
ORDER BY avg_rating DESC