From f40005be0c2b7bbfc9a7d3e5b8976da69ebb8c27 Mon Sep 17 00:00:00 2001 From: Faridun Mirzoev Date: Tue, 7 Apr 2026 10:28:30 -0700 Subject: [PATCH] add AG2 multi-agent RAG example with Weaviate Jupyter notebook and e2e test demonstrating AG2 (formerly AutoGen) multi-agent conversations using Weaviate as the vector store for RAG. - Weaviate Embedded for zero-setup vector search - Automatic embeddings via text2vec-openai module - AG2 GroupChat with Research + Analyst agents - Weaviate search registered as AG2 tool via decorator pattern - Google Colab support --- README.md | 1 + ag2-multiagent-rag/README.md | 53 ++ .../ag2_multiagent_rag_with_weaviate.ipynb | 514 ++++++++++++++++++ ag2-multiagent-rag/test_e2e.py | 269 +++++++++ 4 files changed, 837 insertions(+) create mode 100644 ag2-multiagent-rag/README.md create mode 100644 ag2-multiagent-rag/ag2_multiagent_rag_with_weaviate.ipynb create mode 100644 ag2-multiagent-rag/test_e2e.py diff --git a/README.md b/README.md index b3c38365..95c4b5ad 100644 --- a/README.md +++ b/README.md @@ -36,3 +36,4 @@ Weaviate for cool machine-learning related tasks. | [Workshop Question Answering](question-answering-application-with-weaviate-workshop) | Python notebook | Jupyter/Colab notebook to learn how to get started with Question Answering and Weaviate, given at [Knowledge Graph Conference (KGC) 2022](https://www.knowledgegraph.tech/kgc-2022-tutorial-ml-model-with-the-vector-database-weaviate/) | | [Attendance system using image2vec-neural and own vectors](attendance-system-example) | Python |This example uses image2vec-neural and has an option to use own vectors using OpenCV. It is a smart attendance system example| | [Monitoring Setup with Prometheus & Grafana](monitoring-prometheus-grafana) | yaml | This example does not describe any use case, but rather shows a way of how to start, operate and configure Weaviate with Prometheus-Monitoring and a Grafana Instance with some sample dashboards. | +| [AG2 Multi-Agent RAG with Weaviate](ag2-multiagent-rag) | Python (Jupyter Notebook) | Multi-agent RAG using [AG2](https://ag2.ai/) (formerly AutoGen) with Weaviate as the vector store. AG2 agents collaborate to retrieve and synthesize information via semantic search. | diff --git a/ag2-multiagent-rag/README.md b/ag2-multiagent-rag/README.md new file mode 100644 index 00000000..3972de5d --- /dev/null +++ b/ag2-multiagent-rag/README.md @@ -0,0 +1,53 @@ +# AG2 Multi-Agent RAG with Weaviate + +This example demonstrates how to use [AG2](https://ag2.ai/) (formerly AutoGen) +multi-agent conversations with [Weaviate](https://weaviate.io/) for +Retrieval-Augmented Generation (RAG). + +## Overview + +Two AG2 agents collaborate to answer questions using documents stored in Weaviate: + +- **Research Agent** retrieves relevant documents via Weaviate semantic search +- **Analyst Agent** synthesizes the information into comprehensive answers + +## Prerequisites + +- Python >= 3.10 +- OpenAI API key + +## Quick Start + +```bash +pip install "ag2[openai]>=0.11.4,<1.0" weaviate-client +``` + +Set your API key: + +```bash +export OPENAI_API_KEY="your-api-key" +``` + +Run the notebook: + +```bash +jupyter notebook ag2_multiagent_rag_with_weaviate.ipynb +``` + +Or open directly in Google Colab: + +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/weaviate/weaviate-examples/blob/main/ag2-multiagent-rag/ag2_multiagent_rag_with_weaviate.ipynb) + +## Tech Stack + +- [AG2](https://ag2.ai/) — Multi-agent conversation framework (500K+ monthly PyPI downloads) +- [Weaviate](https://weaviate.io/) — Open-source vector database +- [Weaviate Embedded](https://weaviate.io/developers/weaviate/connections/connect-embedded) — In-process mode (no Docker needed) +- [text2vec-openai](https://weaviate.io/developers/weaviate/model-providers/openai/embeddings) — Automatic embedding generation + +## How It Works + +1. Documents are indexed into Weaviate (auto-embedded via `text2vec-openai`) +2. AG2 Research Agent uses a registered `search_documents` tool to query Weaviate +3. AG2 Analyst Agent synthesizes retrieved information into grounded answers +4. Agents collaborate via AG2 GroupChat with automatic tool execution diff --git a/ag2-multiagent-rag/ag2_multiagent_rag_with_weaviate.ipynb b/ag2-multiagent-rag/ag2_multiagent_rag_with_weaviate.ipynb new file mode 100644 index 00000000..7bf3c2ea --- /dev/null +++ b/ag2-multiagent-rag/ag2_multiagent_rag_with_weaviate.ipynb @@ -0,0 +1,514 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + }, + "colab": { + "name": "ag2_multiagent_rag_with_weaviate.ipynb", + "provenance": [], + "collapsed_sections": [] + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "title-overview" + }, + "source": [ + "# AG2 Multi-Agent RAG with Weaviate\n", + "\n", + "\n", + " \"Open\n", + "\n", + "\n", + "This notebook demonstrates how to use [AG2](https://ag2.ai/) (formerly AutoGen)\n", + "multi-agent conversations with [Weaviate](https://weaviate.io/) as the vector\n", + "database for Retrieval-Augmented Generation (RAG).\n", + "\n", + "**AG2** is a multi-agent conversation framework with 500K+ monthly PyPI downloads,\n", + "4,300+ GitHub stars, and 400+ contributors.\n", + "\n", + "**Weaviate** is an open-source vector database that stores both objects and vectors,\n", + "enabling vector search combined with structured filtering.\n", + "\n", + "In this example, two AG2 agents collaborate:\n", + "- **Research Agent**: Retrieves relevant documents from Weaviate via semantic search\n", + "- **Analyst Agent**: Synthesizes retrieved information into comprehensive answers" + ], + "id": "title-overview" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "install-deps-md" + }, + "source": [ + "## Install Dependencies" + ], + "id": "install-deps-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "install-deps" + }, + "source": [ + "!pip install -U \"ag2[openai]>=0.11.4,<1.0\" weaviate-client" + ], + "id": "install-deps", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "imports-md" + }, + "source": [ + "## Imports and Configuration" + ], + "id": "imports-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "imports-config" + }, + "source": [ + "import os\n", + "\n", + "from autogen import (\n", + " AssistantAgent,\n", + " GroupChat,\n", + " GroupChatManager,\n", + " LLMConfig,\n", + " UserProxyAgent,\n", + ")\n", + "import weaviate\n", + "from weaviate.classes.config import Configure, DataType, Property\n", + "from weaviate.classes.query import MetadataQuery\n", + "\n", + "# Set your OpenAI API key (used for both AG2 and Weaviate text2vec-openai)\n", + "os.environ[\"OPENAI_API_KEY\"] = \"your-api-key\" # Replace with your key" + ], + "id": "imports-config", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "connect-weaviate-md" + }, + "source": [ + "## Connect to Weaviate\n", + "\n", + "We'll use **Weaviate Embedded** which runs in-process — no Docker or server needed.\n", + "For production, use [Weaviate Cloud](https://console.weaviate.cloud/) or\n", + "a self-hosted Docker deployment." + ], + "id": "connect-weaviate-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "connect-weaviate" + }, + "source": [ + "# Connect to Weaviate Embedded (in-process, no Docker needed)\n", + "client = weaviate.connect_to_embedded(\n", + " headers={\"X-OpenAI-Api-Key\": os.environ[\"OPENAI_API_KEY\"]},\n", + ")\n", + "\n", + "print(f\"Weaviate is ready: {client.is_ready()}\")\n", + "\n", + "collection_name = \"AG2_RAG_Demo\"\n", + "\n", + "# Delete collection if exists (for re-running)\n", + "if client.collections.exists(collection_name):\n", + " client.collections.delete(collection_name)\n", + "\n", + "# Create collection with OpenAI vectorizer\n", + "collection = client.collections.create(\n", + " name=collection_name,\n", + " vector_config=Configure.Vectors.text2vec_openai(\n", + " model=\"text-embedding-3-small\",\n", + " ),\n", + " properties=[\n", + " Property(name=\"text\", data_type=DataType.TEXT),\n", + " Property(name=\"source\", data_type=DataType.TEXT),\n", + " ],\n", + ")\n", + "\n", + "print(f\"Collection '{collection_name}' created.\")" + ], + "id": "connect-weaviate", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "prepare-data-md" + }, + "source": [ + "## Prepare and Index Sample Data\n", + "\n", + "We'll index a small set of documents about AI concepts. Weaviate automatically\n", + "generates embeddings using the configured `text2vec-openai` vectorizer — no manual\n", + "embedding step required." + ], + "id": "prepare-data-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "index-documents" + }, + "source": [ + "# Sample documents about AI topics\n", + "documents = [\n", + " {\n", + " \"text\": (\n", + " \"Retrieval-Augmented Generation (RAG) is a technique that combines \"\n", + " \"information retrieval with language model generation. It first retrieves \"\n", + " \"relevant documents from a knowledge base, then uses them as context for \"\n", + " \"generating accurate, grounded responses. RAG reduces hallucination and \"\n", + " \"enables models to access up-to-date information beyond their training data.\"\n", + " ),\n", + " \"source\": \"ai_concepts.md\",\n", + " },\n", + " {\n", + " \"text\": (\n", + " \"Vector databases store data as high-dimensional vectors (embeddings) \"\n", + " \"and enable fast similarity search. Weaviate is an open-source vector database \"\n", + " \"that combines vector search with structured filtering and supports multiple \"\n", + " \"vectorizer modules including OpenAI, Cohere, and Hugging Face transformers.\"\n", + " ),\n", + " \"source\": \"vector_databases.md\",\n", + " },\n", + " {\n", + " \"text\": (\n", + " \"Multi-agent systems use multiple AI agents that collaborate to solve \"\n", + " \"complex tasks. Each agent can have specialized roles, tools, and knowledge. \"\n", + " \"AG2 (formerly AutoGen) is a popular framework for building multi-agent \"\n", + " \"conversations where agents can use tools, write code, and coordinate \"\n", + " \"through structured dialogue patterns.\"\n", + " ),\n", + " \"source\": \"multi_agent_systems.md\",\n", + " },\n", + " {\n", + " \"text\": (\n", + " \"Embedding models convert text into dense numerical vectors that \"\n", + " \"capture semantic meaning. Similar texts produce similar vectors, enabling \"\n", + " \"semantic search. Popular embedding models include OpenAI text-embedding-3-small, \"\n", + " \"sentence-transformers, and BGE models. The choice of embedding model \"\n", + " \"significantly impacts retrieval quality in RAG systems.\"\n", + " ),\n", + " \"source\": \"embeddings.md\",\n", + " },\n", + " {\n", + " \"text\": (\n", + " \"Chunking is the process of splitting documents into smaller pieces \"\n", + " \"for embedding and retrieval. Common strategies include fixed-size chunking, \"\n", + " \"recursive character splitting, and semantic chunking. Optimal chunk size \"\n", + " \"depends on the use case: 256-512 tokens for precise retrieval, 1000+ tokens \"\n", + " \"for broader context. Overlap between chunks helps preserve context.\"\n", + " ),\n", + " \"source\": \"chunking_strategies.md\",\n", + " },\n", + "]\n", + "\n", + "# Insert documents — Weaviate auto-vectorizes via text2vec-openai\n", + "collection = client.collections.get(collection_name)\n", + "\n", + "with collection.batch.dynamic() as batch:\n", + " for doc in documents:\n", + " batch.add_object(properties=doc)\n", + "\n", + "print(\n", + " f\"Indexed {len(documents)} documents into Weaviate collection '{collection_name}'\"\n", + ")" + ], + "id": "index-documents", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "test-search-md" + }, + "source": [ + "## Test Semantic Search\n", + "\n", + "Verify that Weaviate retrieval works before connecting it to AG2 agents." + ], + "id": "test-search-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "test-search" + }, + "source": [ + "# Quick test — Weaviate near_text handles embedding automatically\n", + "results = collection.query.near_text(\n", + " query=\"What is RAG?\",\n", + " limit=3,\n", + " return_metadata=MetadataQuery(distance=True),\n", + ")\n", + "\n", + "for obj in results.objects:\n", + " print(f\"Source: {obj.properties['source']} (distance: {obj.metadata.distance:.4f})\")\n", + " print(f\" {obj.properties['text'][:100]}...\")\n", + " print()" + ], + "id": "test-search", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "search-tool-md" + }, + "source": [ + "## Define Weaviate Search Function for AG2\n", + "\n", + "Create a search function that AG2 agents will use as a registered tool to retrieve\n", + "relevant documents from the Weaviate collection." + ], + "id": "search-tool-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "search-function" + }, + "source": [ + "def search_weaviate(query: str, top_k: int = 3) -> str:\n", + " \"\"\"\n", + " Search Weaviate collection for relevant documents.\n", + "\n", + " Args:\n", + " query: Search query string.\n", + " top_k: Number of results to return.\n", + "\n", + " Returns:\n", + " Formatted string with retrieved documents and sources.\n", + " \"\"\"\n", + " results = collection.query.near_text(\n", + " query=query,\n", + " limit=top_k,\n", + " return_metadata=MetadataQuery(distance=True),\n", + " )\n", + "\n", + " formatted = []\n", + " for i, obj in enumerate(results.objects, 1):\n", + " source = obj.properties[\"source\"]\n", + " text = obj.properties[\"text\"]\n", + " distance = obj.metadata.distance\n", + " formatted.append(f\"[{i}] Source: {source} (distance: {distance:.4f})\\n{text}\")\n", + "\n", + " return (\n", + " \"\\n\\n---\\n\\n\".join(formatted) if formatted else \"No relevant documents found.\"\n", + " )" + ], + "id": "search-function", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ag2-agents-md" + }, + "source": [ + "## Set Up AG2 Multi-Agent System\n", + "\n", + "Create AG2 agents that collaborate to answer questions using Weaviate retrieval:\n", + "\n", + "1. **Research Agent** — Uses the `search_weaviate` tool to find relevant documents\n", + "2. **Analyst Agent** — Synthesizes retrieved information into a final answer\n", + "3. **User Proxy** — Orchestrates the conversation and executes tool calls" + ], + "id": "ag2-agents-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "ag2-setup-run" + }, + "source": [ + "# AG2 LLM Configuration\n", + "llm_config = LLMConfig(\n", + " {\n", + " \"model\": \"gpt-4o-mini\",\n", + " \"api_key\": os.environ[\"OPENAI_API_KEY\"],\n", + " \"api_type\": \"openai\",\n", + " }\n", + ")\n", + "\n", + "# Create agents\n", + "researcher = AssistantAgent(\n", + " name=\"researcher\",\n", + " system_message=(\n", + " \"You are a research agent. When asked a question, use the search_documents \"\n", + " \"tool to retrieve relevant documents from the Weaviate knowledge base. \"\n", + " \"Present your findings clearly with source references. If results are \"\n", + " \"insufficient, try rephrasing your search query.\"\n", + " ),\n", + " llm_config=llm_config,\n", + ")\n", + "\n", + "analyst = AssistantAgent(\n", + " name=\"analyst\",\n", + " system_message=(\n", + " \"You are an analyst. Based on the researcher's findings, synthesize the \"\n", + " \"information into a comprehensive, well-structured answer. Always reference \"\n", + " \"the source documents. End with TERMINATE when done.\"\n", + " ),\n", + " llm_config=llm_config,\n", + ")\n", + "\n", + "user_proxy = UserProxyAgent(\n", + " name=\"user_proxy\",\n", + " human_input_mode=\"NEVER\",\n", + " max_consecutive_auto_reply=10,\n", + " code_execution_config=False,\n", + " is_termination_msg=lambda x: x.get(\"content\", \"\")\n", + " and \"TERMINATE\" in x.get(\"content\", \"\"),\n", + ")\n", + "\n", + "\n", + "# Register Weaviate search as AG2 tool\n", + "@user_proxy.register_for_execution()\n", + "@researcher.register_for_llm(\n", + " description=(\n", + " \"Search the Weaviate vector database for relevant documents. \"\n", + " \"Returns document chunks with similarity distances and source references. \"\n", + " \"Use specific search queries for best results.\"\n", + " )\n", + ")\n", + "def search_documents(query: str, top_k: int = 3) -> str:\n", + " \"\"\"Search Weaviate for relevant documents.\"\"\"\n", + " return search_weaviate(query, top_k)\n", + "\n", + "\n", + "# Set up GroupChat\n", + "group_chat = GroupChat(\n", + " agents=[user_proxy, researcher, analyst],\n", + " messages=[],\n", + " max_round=12,\n", + ")\n", + "\n", + "manager = GroupChatManager(\n", + " groupchat=group_chat,\n", + " llm_config=llm_config,\n", + ")\n", + "\n", + "# Run the multi-agent conversation\n", + "user_proxy.run(\n", + " manager,\n", + " message=(\n", + " \"What is RAG and how does it work with vector databases? \"\n", + " \"Also explain how multi-agent systems can improve RAG quality.\"\n", + " ),\n", + ").process()" + ], + "id": "ag2-setup-run", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "cleanup-md" + }, + "source": [ + "## Cleanup\n", + "\n", + "Close the Weaviate client connection and clean up resources." + ], + "id": "cleanup-md" + }, + { + "cell_type": "code", + "metadata": { + "id": "cleanup" + }, + "source": [ + "# Delete the demo collection\n", + "client.collections.delete(collection_name)\n", + "print(f\"Deleted collection '{collection_name}'\")\n", + "\n", + "# Close the Weaviate client\n", + "client.close()\n", + "print(\"Weaviate connection closed.\")" + ], + "id": "cleanup", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "summary-md" + }, + "source": [ + "## Summary\n", + "\n", + "This notebook demonstrated:\n", + "\n", + "1. **Weaviate** as a vector database with automatic embedding via `text2vec-openai`\n", + "2. **AG2** multi-agent conversations with tool-calling capabilities\n", + "3. **Combining both**: AG2 agents using Weaviate semantic search as a registered tool\n", + " for grounded, citation-backed RAG responses\n", + "\n", + "### Key Components\n", + "\n", + "| Component | Role |\n", + "|-----------|------|\n", + "| **Weaviate Embedded** | In-process vector database (no Docker needed) |\n", + "| **text2vec-openai** | Automatic embedding via OpenAI API |\n", + "| **AG2 UserProxy** | Orchestrates conversation, executes tools |\n", + "| **AG2 Research Agent** | Calls `search_documents` tool |\n", + "| **AG2 Analyst Agent** | Synthesizes findings into final answer |\n", + "\n", + "### Next Steps\n", + "\n", + "- Scale up with more documents and a Weaviate Cloud deployment\n", + "- Add more specialized agents (fact-checker, summarizer, etc.)\n", + "- Use Weaviate hybrid search (vector + BM25) for better retrieval\n", + "- Add metadata filtering for domain-specific queries\n", + "- Explore Weaviate generative modules for built-in RAG\n", + "\n", + "### Resources\n", + "\n", + "- [AG2 Documentation](https://docs.ag2.ai/)\n", + "- [Weaviate Documentation](https://weaviate.io/developers/weaviate)\n", + "- [AG2 GitHub](https://github.com/ag2ai/ag2)\n", + "- [Weaviate GitHub](https://github.com/weaviate/weaviate)" + ], + "id": "summary-md" + } + ] +} \ No newline at end of file diff --git a/ag2-multiagent-rag/test_e2e.py b/ag2-multiagent-rag/test_e2e.py new file mode 100644 index 00000000..f8e84fd7 --- /dev/null +++ b/ag2-multiagent-rag/test_e2e.py @@ -0,0 +1,269 @@ +""" +End-to-end test for the AG2 Multi-Agent RAG with Weaviate example. + +Requires: + - OPENAI_API_KEY environment variable set with a valid key + - pip install "ag2[openai]>=0.11.4,<1.0" weaviate-client + +Usage: + export OPENAI_API_KEY="sk-..." + python test_e2e.py +""" + +import os +import sys + +from autogen import ( + AssistantAgent, + GroupChat, + GroupChatManager, + LLMConfig, + UserProxyAgent, +) +import weaviate +from weaviate.classes.config import Configure, DataType, Property +from weaviate.classes.query import MetadataQuery + + +def main(): + # ------------------------------------------------------------------- + # 1. Validate API key + # ------------------------------------------------------------------- + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key or api_key == "your-api-key": + print("ERROR: Set OPENAI_API_KEY environment variable to a valid key.") + sys.exit(1) + + # ------------------------------------------------------------------- + # 2. Connect to Weaviate Embedded + # ------------------------------------------------------------------- + print("Connecting to Weaviate Embedded...") + client = weaviate.connect_to_embedded( + headers={"X-OpenAI-Api-Key": api_key}, + ) + print(f"Weaviate is ready: {client.is_ready()}") + + collection_name = "AG2_RAG_Demo" + + try: + # --------------------------------------------------------------- + # 3. Create collection + # --------------------------------------------------------------- + if client.collections.exists(collection_name): + client.collections.delete(collection_name) + + collection = client.collections.create( + name=collection_name, + vector_config=Configure.Vectors.text2vec_openai( + model="text-embedding-3-small", + ), + properties=[ + Property(name="text", data_type=DataType.TEXT), + Property(name="source", data_type=DataType.TEXT), + ], + ) + print(f"Collection '{collection_name}' created.") + + # --------------------------------------------------------------- + # 4. Index documents + # --------------------------------------------------------------- + documents = [ + { + "text": ( + "Retrieval-Augmented Generation (RAG) is a technique that combines " + "information retrieval with language model generation. It first retrieves " + "relevant documents from a knowledge base, then uses them as context for " + "generating accurate, grounded responses. RAG reduces hallucination and " + "enables models to access up-to-date information beyond their training data." + ), + "source": "ai_concepts.md", + }, + { + "text": ( + "Vector databases store data as high-dimensional vectors (embeddings) " + "and enable fast similarity search. Weaviate is an open-source vector database " + "that combines vector search with structured filtering and supports multiple " + "vectorizer modules including OpenAI, Cohere, and Hugging Face transformers." + ), + "source": "vector_databases.md", + }, + { + "text": ( + "Multi-agent systems use multiple AI agents that collaborate to solve " + "complex tasks. Each agent can have specialized roles, tools, and knowledge. " + "AG2 (formerly AutoGen) is a popular framework for building multi-agent " + "conversations where agents can use tools, write code, and coordinate " + "through structured dialogue patterns." + ), + "source": "multi_agent_systems.md", + }, + { + "text": ( + "Embedding models convert text into dense numerical vectors that " + "capture semantic meaning. Similar texts produce similar vectors, enabling " + "semantic search. Popular embedding models include OpenAI text-embedding-3-small, " + "sentence-transformers, and BGE models. The choice of embedding model " + "significantly impacts retrieval quality in RAG systems." + ), + "source": "embeddings.md", + }, + { + "text": ( + "Chunking is the process of splitting documents into smaller pieces " + "for embedding and retrieval. Common strategies include fixed-size chunking, " + "recursive character splitting, and semantic chunking. Optimal chunk size " + "depends on the use case: 256-512 tokens for precise retrieval, 1000+ tokens " + "for broader context. Overlap between chunks helps preserve context." + ), + "source": "chunking_strategies.md", + }, + ] + + collection = client.collections.get(collection_name) + + with collection.batch.dynamic() as batch: + for doc in documents: + batch.add_object(properties=doc) + + print( + f"Indexed {len(documents)} documents into Weaviate collection '{collection_name}'" + ) + + # --------------------------------------------------------------- + # 5. Test semantic search + # --------------------------------------------------------------- + print("\n--- Testing Weaviate semantic search ---") + results = collection.query.near_text( + query="What is RAG?", + limit=3, + return_metadata=MetadataQuery(distance=True), + ) + + for obj in results.objects: + print( + f"Source: {obj.properties['source']} (distance: {obj.metadata.distance:.4f})" + ) + print(f" {obj.properties['text'][:100]}...") + print() + + assert len(results.objects) > 0, "Weaviate search returned no results" + print("Weaviate search OK.\n") + + # --------------------------------------------------------------- + # 6. Define search tool + # --------------------------------------------------------------- + def search_weaviate(query: str, top_k: int = 3) -> str: + results = collection.query.near_text( + query=query, + limit=top_k, + return_metadata=MetadataQuery(distance=True), + ) + + formatted = [] + for i, obj in enumerate(results.objects, 1): + source = obj.properties["source"] + text = obj.properties["text"] + distance = obj.metadata.distance + formatted.append( + f"[{i}] Source: {source} (distance: {distance:.4f})\n{text}" + ) + + return ( + "\n\n---\n\n".join(formatted) + if formatted + else "No relevant documents found." + ) + + # --------------------------------------------------------------- + # 7. Set up AG2 agents and run multi-agent conversation + # --------------------------------------------------------------- + print("--- Setting up AG2 agents ---") + + llm_config = LLMConfig( + { + "model": "gpt-4o-mini", + "api_key": api_key, + "api_type": "openai", + } + ) + + researcher = AssistantAgent( + name="researcher", + system_message=( + "You are a research agent. When asked a question, use the search_documents " + "tool to retrieve relevant documents from the Weaviate knowledge base. " + "Present your findings clearly with source references. If results are " + "insufficient, try rephrasing your search query." + ), + llm_config=llm_config, + ) + + analyst = AssistantAgent( + name="analyst", + system_message=( + "You are an analyst. Based on the researcher's findings, synthesize the " + "information into a comprehensive, well-structured answer. Always reference " + "the source documents. End with TERMINATE when done." + ), + llm_config=llm_config, + ) + + user_proxy = UserProxyAgent( + name="user_proxy", + human_input_mode="NEVER", + max_consecutive_auto_reply=10, + code_execution_config=False, + is_termination_msg=lambda x: x.get("content", "") + and "TERMINATE" in x.get("content", ""), + ) + + @user_proxy.register_for_execution() + @researcher.register_for_llm( + description=( + "Search the Weaviate vector database for relevant documents. " + "Returns document chunks with similarity distances and source references. " + "Use specific search queries for best results." + ) + ) + def search_documents(query: str, top_k: int = 3) -> str: + """Search Weaviate for relevant documents.""" + return search_weaviate(query, top_k) + + group_chat = GroupChat( + agents=[user_proxy, researcher, analyst], + messages=[], + max_round=12, + ) + + manager = GroupChatManager( + groupchat=group_chat, + llm_config=llm_config, + ) + + print("--- Running multi-agent conversation (real LLM calls) ---\n") + + user_proxy.run( + manager, + message=( + "What is RAG and how does it work with vector databases? " + "Also explain how multi-agent systems can improve RAG quality." + ), + ).process() + + print("\n--- Multi-agent conversation completed ---") + + finally: + # --------------------------------------------------------------- + # 8. Cleanup + # --------------------------------------------------------------- + if client.collections.exists(collection_name): + client.collections.delete(collection_name) + print(f"Deleted collection '{collection_name}'") + client.close() + print("Weaviate connection closed.") + + print("\nAll done. Test passed.") + + +if __name__ == "__main__": + main()