Compartir a través de

AI Agent with Vector Search Sample using Go and Azure DocumentDB

Browse code

This sample demonstrates building an intelligent AI agent using Go with Azure DocumentDB vector search capabilities. The agent provides hotel recommendations based on user queries by leveraging a two-agent architecture: a planner agent for query refinement and vector search, and a synthesizer agent for comparative analysis of results.

⚠️ Important Note on agentic architecture
This sample is a reference implementation demonstrating agentic patterns in Go. It uses a custom-built agent architecture rather than an agent framework, which is the recommended approach for production agentic applications.

Architecture

The sample implements a two-agent architecture for hotel recommendations:

  • Planner Agent (gpt-4o-mini): Refines queries and executes vector search using a custom tool
  • Synthesizer Agent (gpt-4o): Analyzes search results and provides comparative recommendations

The application uses:

  • Embeddings: text-embedding-3-small (1536 dimensions)
  • Vector Database: Azure DocumentDB
  • Vector Algorithms: Supports IVF, HNSW, and DiskANN

Architecture Flow

flowchart LR
    subgraph Input
        U[User Query]
    end

    subgraph Planner["Planner Agent (gpt-4o-mini)"]
        P1[Refine user query]
        P2[OpenAI function calling]
        P3[Return structured data<br/>with similarity scores]
    end

    subgraph Tool["Vector Search Tool"]
        T1[Generate embeddings<br/>text-embedding-3-small]
        T2[Execute MongoDB vector search<br/>cosine similarity]
        T3[Format results]
    end

    subgraph Synthesizer["Synthesizer Agent (gpt-4o)"]
        S1[Analyze top 3 results<br/>comparatively]
        S2[Generate natural language<br/>recommendations]
        S3[Explain tradeoffs<br/>between options]
    end

    subgraph Output
        R[Hotel Recommendations]
    end

    U --> P1 --> P2 --> T1 --> T2 --> T3 --> P3 --> S1 --> S2 --> S3 --> R

Implementation Notes:

  • This is a custom implementation using OpenAI SDK directly with no framework
  • Uses OpenAI function calling for tool integration
  • Linear workflow: Planner → Tool → Synthesizer
  • Stateless execution with no conversation history
  • Suitable for single-turn query/response scenarios

Prerequisites

Azure Resources

  1. Azure OpenAI Service with deployed models:

    • text-embedding-3-small deployment (Embeddings)
    • gpt-4o-mini deployment (Planner Agent)
    • gpt-4o deployment (Synthesizer Agent)

  2. Azure DocumentDB cluster with:

    • Vector search support enabled
    • Appropriate cluster tier (M10+ for IVF, M30+ for HNSW/DiskANN)
    • Firewall rules configured to allow your IP. For more information, see Grant access from your IP address.

Development Tools

  • Go 1.21 or later

Project Structure

vector-search-agent-go/
├── cmd/
│   ├── agent/          # Main agent application
│   ├── upload/         # Data upload utility
│   └── cleanup/        # Database cleanup utility
├── internal/
│   ├── models/         # Hotel data models
│   ├── clients/        # Azure OpenAI client
│   ├── vectorstore/    # Azure DocumentDB vector store operations
│   ├── agents/         # Agent and tool implementations
│   │   ├── agents.go   # Planner and synthesizer agents
│   │   └── tools.go    # Vector search tool definition
│   └── prompts/        # System prompts and tool definitions
├── go.mod
├── go.sum
└── README.md

Installation

  1. Clone the repository:
git clone https://github.com/Azure-Samples/documentdb-samples.git
cd documentdb-samples/vector-search-agent-go
  1. Install dependencies:
go mod download
  1. Configure authentication by creating a .env file.

The sample uses a two-agent architecture (planner + synthesizer) with three model deployments (two chat models + embeddings). The environment variables are configured for each model deployment.

  • AZURE_OPENAI_PLANNER_DEPLOYMENT: Your gpt-4o-mini deployment name
  • AZURE_OPENAI_SYNTH_DEPLOYMENT: Your gpt-4o deployment name
  • AZURE_OPENAI_EMBEDDING_DEPLOYMENT: Your text-embedding-3-small deployment name

You can choose between two authentication methods: passwordless authentication using Azure Identity (recommended) or traditional connection string and API key.

Use passwordless authentication with both Azure OpenAI and Azure DocumentDB. Set USE_PASSWORDLESS=true, AZURE_OPENAI_ENDPOINT, and AZURE_DOCUMENTDB_CLUSTER.

# Enable passwordless authentication
USE_PASSWORDLESS=true

# Azure OpenAI Configuration (passwordless)
AZURE_OPENAI_ENDPOINT=https://your-openai-instance.openai.azure.com/

# Azure DocumentDB (passwordless)
AZURE_DOCUMENTDB_CLUSTER=your-mongo-cluster-name
AZURE_DOCUMENTDB_DATABASENAME=Hotels
AZURE_DOCUMENTDB_COLLECTION=hotel_data
AZURE_DOCUMENTDB_INDEX_NAME=vectorIndex

# Optional settings
DEBUG=true
QUERY=quintessential lodging near running trails, eateries, retail
NEAREST_NEIGHBORS=5

Prerequisites for passwordless authentication:

  • Ensure you're logged in to Azure: az login
  • OR have appropriate managed identity/service principal/workload identity configured
  • Grant your identity the following roles:
    • Cognitive Services OpenAI User on the Azure OpenAI resource
    • DocumentDB Account Contributor and Cosmos DB Account Reader Role on the Azure DocumentDB resource

Option 2: Connection String and API Key Authentication

Use key-based authentication by setting USE_PASSWORDLESS=false and providing AZURE_OPENAI_API_KEY and AZURE_DOCUMENTDB_CONNECTION_STRING in your .env file.

# Disable passwordless authentication
USE_PASSWORDLESS=false

# Azure OpenAI Configuration (API key)
AZURE_OPENAI_ENDPOINT=https://your-openai-instance.openai.azure.com/
AZURE_OPENAI_API_KEY=your-azure-openai-api-key

# Azure DocumentDB (connection string)
AZURE_DOCUMENTDB_CONNECTION_STRING=mongodb+srv://username:password@cluster.mongocluster.cosmos.azure.com/
AZURE_DOCUMENTDB_DATABASENAME=vectorSearchDB
AZURE_DOCUMENTDB_COLLECTION=hotel_data
AZURE_DOCUMENTDB_INDEX_NAME=vectorIndex

Usage

Upload the hotel data with embeddings:

go run cmd/upload/main.go

The upload command will:

  • Load hotels from the JSON file
  • Generate embeddings for each hotel
  • Insert documents into Azure DocumentDB
  • Create a vector index

2. Run the Agent

Run the hotel recommendation agent:

go run cmd/agent/main.go

The agent will:

  • Connect to the existing vector store
  • Execute the planner agent (query refinement + vector search)
  • Execute the synthesizer agent (comparative analysis)
  • Display the final recommendation

Example output:

Query: quintessential lodging near running trails, eateries, retail
Nearest Neighbors: 5

--- PLANNER ---
Tool: search_hotels_collection
Query: quintessential lodging near running trails, eateries, and retail shops with good amenities and comfort
K: 5
Hotel #1: Country Comfort Inn, Score: 0.513081
Hotel #2: Nordick's Valley Motel, Score: 0.494860
Hotel #3: Trails End Motel, Score: 0.492937
Hotel #4: Lakefront Captain Inn, Score: 0.478377
Hotel #5: White Mountain Lodge & Suites, Score: 0.477240

--- SYNTHESIZER ---
Context size: 3233 characters

--- FINAL ANSWER ---
1. COMPARISON SUMMARY:
• Country Comfort Inn has the lowest rating (2.5) but offers free parking, free wifi, and pet-friendly extended-stay suites. It is located in Bellevue, WA near a lake.
• Nordick's Valley Motel has the highest rating (4.5) and offers free parking and continental breakfast. It is a boutique hotel near hiking and wine tasting in Washington D.C.
• Trails End Motel has a mid-range rating (3.2) and offers free wifi, free parking, and a hot breakfast buffet. It is a budget hotel located 8 miles from Downtown Scottsdale, AZ.

Key tradeoffs:
- Nordick's Valley Motel excels in rating but lacks free wifi and is not near running trails.
- Country Comfort Inn is pet-friendly and offers functional kitchens but has the lowest rating.
- Trails End Motel balances amenities (wifi, breakfast, parking) with a moderate rating but is farther from natural trails.

2. BEST OVERALL:
Trails End Motel is the best overall for its balanced rating, free wifi, free parking, and proximity to Downtown Scottsdale.

3. ALTERNATIVE PICKS:
• Choose Nordick's Valley Motel if high rating and boutique charm near hiking or wine tasting are priorities.
• Choose Country Comfort Inn if pet-friendly extended stays near a lake are essential.

3. Cleanup

To delete the test database:

go run cmd/cleanup/main.go

Key Implementation Details

No Framework

This implementation demonstrates how to build an agent system without using a framework:

  • Direct Azure OpenAI SDK usage: Uses the official github.com/openai/openai-go/v3 package with Azure support
  • Passwordless authentication: Leverages github.com/Azure/azure-sdk-for-go/sdk/azidentity for secure, keyless authentication
  • Custom tool calling: Implements tool definitions and execution manually
  • Manual agent orchestration: Explicit control flow between planner and synthesizer
  • MongoDB native driver: Uses go.mongodb.org/mongo-driver directly with OIDC authentication support

Two-Agent Pattern

Planner Agent:

  • Uses OpenAI function calling to invoke the vector search tool
  • Refines user queries for better semantic search
  • Returns structured hotel data with similarity scores

Synthesizer Agent:

  • Analyzes top 3 results comparatively
  • Generates natural language recommendations
  • Explains tradeoffs between options

Vector Search Tool

The search tool:

  1. Generates embeddings using text-embedding-3-small
  2. Executes MongoDB vector search with cosine similarity
  3. Formats results for the synthesizer agent

Configuration Options

Authentication Methods

The application supports two authentication methods for both Azure OpenAI and Azure DocumentDB:

  1. Passwordless Authentication (Recommended):

    • Set USE_PASSWORDLESS=true in your .env file
    • Uses Azure Identity (DefaultAzureCredential) for authentication
    • For Azure OpenAI: Omit AZURE_OPENAI_API_KEY and provide only AZURE_OPENAI_ENDPOINT
    • For Azure DocumentDB: Omit AZURE_DOCUMENTDB_CONNECTION_STRING and provide AZURE_DOCUMENTDB_CLUSTER instead
    • Supports multiple credential types: Azure CLI, Managed Identity, Service Principal, Workload Identity
    • More secure as no secrets are stored in code or config files
    • Requires appropriate RBAC roles assigned to your identity

  2. API Key / Connection String Authentication:

    • Set USE_PASSWORDLESS=false in your .env file (or omit the variable)
    • For Azure OpenAI: Provide AZURE_OPENAI_API_KEY
    • For Azure DocumentDB: Provide AZURE_DOCUMENTDB_CONNECTION_STRING
    • Simpler setup but requires managing secrets

Auto-detection: If USE_PASSWORDLESS is not explicitly set, the application automatically detects which authentication method to use:

  • Uses passwordless for OpenAI if AZURE_OPENAI_API_KEY is not set
  • Uses passwordless for Azure DocumentDB if AZURE_DOCUMENTDB_CONNECTION_STRING is not set but AZURE_DOCUMENTDB_CLUSTER is set

Vector Index Algorithms

  • IVF (default): VECTOR_INDEX_ALGORITHM=vector-ivf
  • HNSW: VECTOR_INDEX_ALGORITHM=vector-hnsw
  • DiskANN: VECTOR_INDEX_ALGORITHM=vector-diskann

Similarity Metrics

  • Cosine (default): VECTOR_SIMILARITY=COS
  • Euclidean: VECTOR_SIMILARITY=L2
  • Inner Product: VECTOR_SIMILARITY=IP

Debug Mode

Enable detailed logging:

DEBUG=true

Troubleshooting

Connection Errors

If you get connection timeouts:

  • Verify your IP is in the Azure DocumentDB firewall rules
  • Check your connection string is correct
  • Ensure the cluster is running

Rate Limiting

If you encounter 429 errors:

  • Increase TPM quotas in Azure portal
  • Reduce the number of hotels processed
  • Add retry logic with exponential backoff

Tool Not Called

If the planner doesn't call the search tool:

  • Verify the tool definition matches the prompt
  • Check the planner system prompt is explicit
  • Enable DEBUG mode to see LLM responses