diff --git a/.github/workflows/contrib-tests.yml b/.github/workflows/contrib-tests.yml index ac0240eb43..40aa06303e 100644 --- a/.github/workflows/contrib-tests.yml +++ b/.github/workflows/contrib-tests.yml @@ -87,6 +87,17 @@ jobs: --health-retries 5 ports: - 5432:5432 + couchbase: + image: couchbase:enterprise-7.6.3 + ports: + - "8091-8095:8091-8095" + - "11210:11210" + - "9102:9102" + healthcheck: # checks couchbase server is up + test: ["CMD", "curl", "-v", "http://localhost:8091/pools"] + interval: 20s + timeout: 20s + retries: 5 mongodb: image: mongodb/mongodb-atlas-local:latest ports: @@ -111,6 +122,9 @@ jobs: - name: Install mongodb when on linux run: | pip install -e .[retrievechat-mongodb] + - name: Install couchbase when on linux + run: | + pip install -e .[retrievechat-couchbase] - name: Install unstructured when python-version is 3.9 and on linux if: matrix.python-version == '3.9' run: | diff --git a/autogen/agentchat/contrib/vectordb/base.py b/autogen/agentchat/contrib/vectordb/base.py index 1454d65318..ad0c4d5250 100644 --- a/autogen/agentchat/contrib/vectordb/base.py +++ b/autogen/agentchat/contrib/vectordb/base.py @@ -207,7 +207,7 @@ class VectorDBFactory: Factory class for creating vector databases. """ - PREDEFINED_VECTOR_DB = ["chroma", "pgvector", "mongodb", "qdrant"] + PREDEFINED_VECTOR_DB = ["chroma", "pgvector", "mongodb", "qdrant", "couchbase"] @staticmethod def create_vector_db(db_type: str, **kwargs) -> VectorDB: @@ -237,6 +237,10 @@ def create_vector_db(db_type: str, **kwargs) -> VectorDB: from .qdrant import QdrantVectorDB return QdrantVectorDB(**kwargs) + if db_type.lower() in ["couchbase", "couchbasedb", "capella"]: + from .couchbase import CouchbaseVectorDB + + return CouchbaseVectorDB(**kwargs) else: raise ValueError( f"Unsupported vector database type: {db_type}. Valid types are {VectorDBFactory.PREDEFINED_VECTOR_DB}." diff --git a/autogen/agentchat/contrib/vectordb/couchbase.py b/autogen/agentchat/contrib/vectordb/couchbase.py new file mode 100644 index 0000000000..3052a8a9f1 --- /dev/null +++ b/autogen/agentchat/contrib/vectordb/couchbase.py @@ -0,0 +1,386 @@ +import json +import time +from datetime import timedelta +from typing import Any, Callable, Dict, List, Literal, Tuple, Union + +import numpy as np +from couchbase import search +from couchbase.auth import PasswordAuthenticator +from couchbase.cluster import Cluster, ClusterOptions +from couchbase.collection import Collection +from couchbase.management.search import SearchIndex +from couchbase.options import SearchOptions +from couchbase.vector_search import VectorQuery, VectorSearch +from sentence_transformers import SentenceTransformer + +from .base import Document, ItemID, QueryResults, VectorDB +from .utils import get_logger + +logger = get_logger(__name__) + +DEFAULT_BATCH_SIZE = 1000 +_SAMPLE_SENTENCE = ["The weather is lovely today in paradise."] +TEXT_KEY = "content" +EMBEDDING_KEY = "embedding" + + +class CouchbaseVectorDB(VectorDB): + """ + A vector database implementation that uses Couchbase as the backend. + """ + + def __init__( + self, + connection_string: str = "couchbase://localhost", + username: str = "Administrator", + password: str = "password", + bucket_name: str = "vector_db", + embedding_function: Callable = SentenceTransformer("all-MiniLM-L6-v2").encode, + scope_name: str = "_default", + collection_name: str = "_default", + index_name: str = None, + ): + """ + Initialize the vector database. + Args: + connection_string (str): The Couchbase connection string to connect to. Default is 'couchbase://localhost'. + username (str): The username for Couchbase authentication. Default is 'Administrator'. + password (str): The password for Couchbase authentication. Default is 'password'. + bucket_name (str): The name of the bucket. Default is 'vector_db'. + embedding_function (Callable): The embedding function used to generate the vector representation. Default is SentenceTransformer("all-MiniLM-L6-v2").encode. + scope_name (str): The name of the scope. Default is '_default'. + collection_name (str): The name of the collection to create for this vector database. Default is '_default'. + index_name (str): Index name for the vector database. Default is None. + overwrite (bool): Whether to overwrite existing data. Default is False. + wait_until_index_ready (float | None): Blocking call to wait until the database indexes are ready. None means no wait. Default is None. + wait_until_document_ready (float | None): Blocking call to wait until the database documents are ready. None means no wait. Default is None. + """ + + self.embedding_function = embedding_function + self.index_name = index_name + + # This will get the model dimension size by computing the embeddings dimensions + self.dimensions = self._get_embedding_size() + + try: + auth = PasswordAuthenticator(username, password) + cluster = Cluster(connection_string, ClusterOptions(auth)) + cluster.wait_until_ready(timedelta(seconds=5)) + self.cluster = cluster + + self.bucket = cluster.bucket(bucket_name) + self.scope = self.bucket.scope(scope_name) + self.collection = self.scope.collection(collection_name) + self.active_collection = self.collection + + logger.debug("Successfully connected to Couchbase") + except Exception as err: + raise ConnectionError("Could not connect to Couchbase server") from err + + def search_index_exists(self, index_name: str): + """Check if the specified index is ready""" + try: + search_index_mgr = self.scope.search_indexes() + index = search_index_mgr.get_index(index_name) + return index.is_valid() + except Exception: + return False + + def _get_embedding_size(self): + return len(self.embedding_function(_SAMPLE_SENTENCE)[0]) + + def create_collection( + self, + collection_name: str, + overwrite: bool = False, + get_or_create: bool = True, + ) -> Collection: + """ + Create a collection in the vector database and create a vector search index in the collection. + Args: + collection_name: str | The name of the collection. + overwrite: bool | Whether to overwrite the collection if it exists. Default is False. + get_or_create: bool | Whether to get or create the collection. Default is True + """ + if overwrite: + self.delete_collection(collection_name) + + try: + collection_mgr = self.bucket.collections() + collection_mgr.create_collection(self.scope.name, collection_name) + self.cluster.query(f"CREATE PRIMARY INDEX ON {self.bucket.name}.{self.scope.name}.{collection_name}") + + except Exception: + if not get_or_create: + raise ValueError(f"Collection {collection_name} already exists.") + else: + logger.debug(f"Collection {collection_name} already exists. Getting the collection.") + + collection = self.scope.collection(collection_name) + self.create_index_if_not_exists(index_name=self.index_name, collection=collection) + return collection + + def create_index_if_not_exists(self, index_name: str = "vector_index", collection=None) -> None: + """ + Creates a vector search index on the specified collection in Couchbase. + Args: + index_name (str, optional): The name of the vector search index to create. Defaults to "vector_search_index". + collection (Collection, optional): The Couchbase collection to create the index on. Defaults to None. + """ + if not self.search_index_exists(index_name): + self.create_vector_search_index(collection, index_name) + + def get_collection(self, collection_name: str|None = None) -> Collection: + """ + Get the collection from the vector database. + Args: + collection_name: str | The name of the collection. Default is None. If None, return the + current active collection. + Returns: + Collection | The collection object. + """ + if collection_name is None: + if self.active_collection is None: + raise ValueError("No collection is specified.") + else: + logger.debug( + f"No collection is specified. Using current active collection {self.active_collection.name}." + ) + else: + self.active_collection = self.scope.collection(collection_name) + + return self.active_collection + + def delete_collection(self, collection_name: str) -> None: + """ + Delete the collection from the vector database. + Args: + collection_name: str | The name of the collection. + """ + try: + collection_mgr = self.bucket.collections() + collection_mgr.drop_collection(self.scope.name, collection_name) + except Exception as e: + logger.error(f"Error deleting collection: {e}") + + def create_vector_search_index( + self, + collection, + index_name: Union[str, None] = "vector_index", + similarity: Literal["l2_norm", "dot_product"] = "dot_product", + ) -> None: + """Create a vector search index in the collection.""" + search_index_mgr = self.scope.search_indexes() + dims = self._get_embedding_size() + index_definition = { + "type": "fulltext-index", + "name": index_name, + "sourceType": "couchbase", + "sourceName": self.bucket.name, + "planParams": {"maxPartitionsPerPIndex": 1024, "indexPartitions": 1}, + "params": { + "doc_config": { + "docid_prefix_delim": "", + "docid_regexp": "", + "mode": "scope.collection.type_field", + "type_field": "type", + }, + "mapping": { + "analysis": {}, + "default_analyzer": "standard", + "default_datetime_parser": "dateTimeOptional", + "default_field": "_all", + "default_mapping": {"dynamic": True, "enabled": False}, + "default_type": "_default", + "docvalues_dynamic": False, + "index_dynamic": True, + "store_dynamic": True, + "type_field": "_type", + "types": { + f"{self.scope.name}.{collection.name}": { + "dynamic": False, + "enabled": True, + "properties": { + "embedding": { + "dynamic": False, + "enabled": True, + "fields": [ + { + "dims": dims, + "index": True, + "name": "embedding", + "similarity": similarity, + "type": "vector", + "vector_index_optimized_for": "recall", + } + ], + }, + "metadata": {"dynamic": True, "enabled": True}, + "content": { + "dynamic": False, + "enabled": True, + "fields": [ + { + "include_in_all": True, + "index": True, + "name": "content", + "store": True, + "type": "text", + } + ], + }, + }, + } + }, + }, + "store": {"indexType": "scorch", "segmentVersion": 16}, + }, + "sourceParams": {}, + } + + search_index_def = SearchIndex.from_json(json.dumps(index_definition)) + max_attempts = 10 + attempt = 0 + while attempt < max_attempts: + try: + search_index_mgr.upsert_index(search_index_def) + break + except Exception as e: + logger.debug(f"Attempt {attempt + 1}/{max_attempts}: Error creating search index: {e}") + time.sleep(3) + attempt += 1 + + if attempt == max_attempts: + logger.error(f"Error creating search index after {max_attempts} attempts.") + raise RuntimeError(f"Error creating search index after {max_attempts} attempts.") + + logger.info(f"Search index {index_name} created successfully.") + + def upsert_docs( + self, docs: List[Document], collection: Collection, batch_size=DEFAULT_BATCH_SIZE, **kwargs: Any + ) -> None: + if docs[0].get("content") is None: + raise ValueError("The document content is required.") + if docs[0].get("id") is None: + raise ValueError("The document id is required.") + + for i in range(0, len(docs), batch_size): + batch = docs[i : i + batch_size] + docs_to_upsert = dict() + for doc in batch: + doc_id = doc["id"] + embedding = self.embedding_function( + [doc["content"]] + ).tolist() # Gets new embedding even in case of document update + doc_content = { + TEXT_KEY: doc["content"], + "metadata": doc.get("metadata", {}), + EMBEDDING_KEY: embedding, + "id": doc_id, + } + docs_to_upsert[doc_id] = doc_content + collection.upsert_multi(docs_to_upsert) + + def insert_docs( + self, + docs: List[Document], + collection_name: str = None, + upsert: bool = False, + batch_size=DEFAULT_BATCH_SIZE, + **kwargs, + ) -> None: + """Insert Documents and Vector Embeddings into the collection of the vector database. Documents are upserted in all cases.""" + if not docs: + logger.info("No documents to insert.") + return + + collection = self.get_collection(collection_name) + self.upsert_docs(docs, collection, batch_size=batch_size) + + def update_docs( + self, docs: List[Document], collection_name: str = None, batch_size=DEFAULT_BATCH_SIZE, **kwargs: Any + ) -> None: + """Update documents, including their embeddings, in the Collection.""" + collection = self.get_collection(collection_name) + self.upsert_docs(docs, collection, batch_size) + + def delete_docs(self, ids: List[ItemID], collection_name: str = None, batch_size=DEFAULT_BATCH_SIZE, **kwargs): + """Delete documents from the collection of the vector database.""" + collection = self.get_collection(collection_name) + # based on batch size, delete the documents + for i in range(0, len(ids), batch_size): + batch = ids[i : i + batch_size] + collection.remove_multi(batch) + + def get_docs_by_ids( + self, ids: List[ItemID] | None = None, collection_name: str = None, include: List[str] | None = None, **kwargs + ) -> List[Document]: + """Retrieve documents from the collection of the vector database based on the ids.""" + if include is None: + include = [TEXT_KEY, "metadata", "id"] + elif "id" not in include: + include.append("id") + + collection = self.get_collection(collection_name) + if ids is not None: + docs = [collection.get(doc_id) for doc_id in ids] + else: + # Get all documents using couchbase query + include_str = ", ".join(include) + query = f"SELECT {include_str} FROM {self.bucket.name}.{self.scope.name}.{collection.name}" + result = self.cluster.query(query) + docs = [] + for row in result: + docs.append(row) + + return [{k: v for k, v in doc.items() if k in include or k == "id"} for doc in docs] + + def retrieve_docs( + self, + queries: List[str], + collection_name: str = None, + n_results: int = 10, + distance_threshold: float = -1, + **kwargs, + ) -> QueryResults: + """Retrieve documents from the collection of the vector database based on the queries. + Note: Distance threshold is not supported in Couchbase FTS. + """ + + results: QueryResults = [] + for query_text in queries: + query_vector = np.array(self.embedding_function([query_text])).tolist()[0] + query_result = self._vector_search( + query_vector, + n_results, + **kwargs, + ) + results.append(query_result) + return results + + def _vector_search(self, embedding_vector: List[float], n_results: int = 10, **kwargs) -> List[Tuple[Dict, float]]: + """Core vector search using Couchbase FTS.""" + + search_req = search.SearchRequest.create( + VectorSearch.from_vector_query( + VectorQuery( + EMBEDDING_KEY, + embedding_vector, + n_results, + ) + ) + ) + + search_options = SearchOptions(limit=n_results, fields=["*"]) + result = self.scope.search(self.index_name, search_req, search_options) + + docs_with_score = [] + + for row in result.rows(): + doc = row.fields + doc["id"] = row.id + score = row.score + + docs_with_score.append((doc, score)) + + return docs_with_score \ No newline at end of file diff --git a/notebook/agentchat_RetrieveChat_couchbase.ipynb b/notebook/agentchat_RetrieveChat_couchbase.ipynb new file mode 100644 index 0000000000..7e7abbec8f --- /dev/null +++ b/notebook/agentchat_RetrieveChat_couchbase.ipynb @@ -0,0 +1,579 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Using RetrieveChat Powered by Couchbase Capella for Retrieve Augmented Code Generation and Question Answering\n", + "\n", + "AutoGen offers conversable agents powered by LLM, tool or human, which can be used to perform tasks collectively via automated chat. This framework allows tool use and human participation through multi-agent conversation.\n", + "Please find documentation about this feature [here](https://microsoft.github.io/autogen/docs/Use-Cases/agent_chat).\n", + "\n", + "RetrieveChat is a conversational system for retrieval-augmented code generation and question answering. In this notebook, we demonstrate how to utilize RetrieveChat to generate code and answer questions based on customized documentations that are not present in the LLM's training dataset. RetrieveChat uses the `AssistantAgent` and `RetrieveUserProxyAgent`, which is similar to the usage of `AssistantAgent` and `UserProxyAgent` in other notebooks (e.g., [Automated Task Solving with Code Generation, Execution & Debugging](https://github.com/microsoft/autogen/blob/main/notebook/agentchat_auto_feedback_from_code_execution.ipynb)). Essentially, `RetrieveUserProxyAgent` implement a different auto-reply mechanism corresponding to the RetrieveChat prompts.\n", + "\n", + "## Table of Contents\n", + "We'll demonstrate six examples of using RetrieveChat for code generation and question answering:\n", + "\n", + "- [Example 1: Generate code based off docstrings w/o human feedback](#example-1)\n", + "\n", + "````{=mdx}\n", + ":::info Requirements\n", + "Some extra dependencies are needed for this notebook, which can be installed via pip:\n", + "\n", + "```bash\n", + "pip install pyautogen[retrievechat-couchbase] flaml[automl]\n", + "```\n", + "\n", + "For more information, please refer to the [installation guide](/docs/installation/).\n", + ":::\n", + "````\n", + "\n", + "Ensure you have a Couchbase Capella cluster running. Read more on how to get started [here](https://docs.couchbase.com/cloud/get-started/intro.html)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set your API Endpoint\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "models to use: ['gpt-4o-mini']\n" + ] + } + ], + "source": [ + "import os\n", + "import sys\n", + "\n", + "from autogen import AssistantAgent\n", + "\n", + "sys.path.append(os.path.abspath(\"/workspaces/autogen/autogen/agentchat/contrib\"))\n", + "\n", + "from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent\n", + "\n", + "# Accepted file formats for that can be stored in\n", + "# a vector database instance\n", + "from autogen.retrieve_utils import TEXT_FORMATS\n", + "\n", + "config_list = [{\"model\": \"gpt-4o-mini\", \"api_key\": os.environ[\"OPENAI_API_KEY\"], \"api_type\": \"openai\"}]\n", + "assert len(config_list) > 0\n", + "print(\"models to use: \", [config_list[i][\"model\"] for i in range(len(config_list))])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "````{=mdx}\n", + ":::tip\n", + "Learn more about configuring LLMs for agents [here](/docs/topics/llm_configuration).\n", + ":::\n", + "````\n", + "\n", + "## Construct agents for RetrieveChat\n", + "\n", + "We start by initializing the `AssistantAgent` and `RetrieveUserProxyAgent`. The system message needs to be set to \"You are a helpful assistant.\" for AssistantAgent. The detailed instructions are given in the user message. Later we will use the `RetrieveUserProxyAgent.message_generator` to combine the instructions and a retrieval augmented generation task for an initial prompt to be sent to the LLM assistant." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Accepted file formats for `docs_path`:\n", + "['txt', 'json', 'csv', 'tsv', 'md', 'html', 'htm', 'rtf', 'rst', 'jsonl', 'log', 'xml', 'yaml', 'yml', 'pdf']\n" + ] + } + ], + "source": [ + "print(\"Accepted file formats for `docs_path`:\")\n", + "print(TEXT_FORMATS)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 1. create an AssistantAgent instance named \"assistant\"\n", + "assistant = AssistantAgent(\n", + " name=\"assistant\",\n", + " system_message=\"You are a helpful assistant.\",\n", + " llm_config={\n", + " \"timeout\": 600,\n", + " \"cache_seed\": 42,\n", + " \"config_list\": config_list,\n", + " },\n", + ")\n", + "\n", + "# 2. create the RetrieveUserProxyAgent instance named \"ragproxyagent\"\n", + "# Refer to https://microsoft.github.io/autogen/docs/reference/agentchat/contrib/retrieve_user_proxy_agent\n", + "# and https://microsoft.github.io/autogen/docs/reference/agentchat/contrib/vectordb/couchbase\n", + "# for more information on the RetrieveUserProxyAgent and CouchbaseVectorDB\n", + "ragproxyagent = RetrieveUserProxyAgent(\n", + " name=\"ragproxyagent\",\n", + " human_input_mode=\"NEVER\",\n", + " max_consecutive_auto_reply=3,\n", + " retrieve_config={\n", + " \"task\": \"code\",\n", + " \"docs_path\": [\n", + " \"https://raw.githubusercontent.com/microsoft/FLAML/main/website/docs/Examples/Integrate%20-%20Spark.md\",\n", + " \"https://raw.githubusercontent.com/microsoft/FLAML/main/website/docs/Research.md\",\n", + " ],\n", + " \"chunk_token_size\": 2000,\n", + " \"model\": config_list[0][\"model\"],\n", + " \"vector_db\": \"couchbase\", # Couchbase Capella VectorDB\n", + " \"collection_name\": \"demo_collection\", # Couchbase Capella collection name to be utilized/created\n", + " \"db_config\": {\n", + " \"connection_string\": os.environ[\"CB_CONN_STR\"], # Couchbase Capella connection string\n", + " \"username\": os.environ[\"CB_USERNAME\"], # Couchbase Capella username\n", + " \"password\": os.environ[\"CB_PASSWORD\"], # Couchbase Capella password\n", + " \"bucket_name\": \"test_db\", # Couchbase Capella bucket name\n", + " \"scope_name\": \"test_scope\", # Couchbase Capella scope name\n", + " \"index_name\": \"vector_index\", # Couchbase Capella index name to be created\n", + " },\n", + " \"get_or_create\": True, # set to False if you don't want to reuse an existing collection\n", + " \"overwrite\": False, # set to True if you want to overwrite an existing collection, each overwrite will force a index creation and reupload of documents\n", + " },\n", + " code_execution_config=False, # set to False if you don't want to execute the code\n", + ")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Example 1\n", + "\n", + "[Back to top](#table-of-contents)\n", + "\n", + "Use RetrieveChat to help generate sample code and automatically run the code and fix errors if there is any.\n", + "\n", + "Problem: Which API should I use if I want to use FLAML for a classification task and I want to train the model in 30 seconds. Use spark to parallel the training. Force cancel jobs if time limit is reached.\n", + "\n", + "Note: You may need to create an index on the cluster to query" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2024-10-16 12:08:07,062 - autogen.agentchat.contrib.retrieve_user_proxy_agent - INFO - \u001b[32mUse the existing collection `demo_collection`.\u001b[0m\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Trying to create collection.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2024-10-16 12:08:07,953 - autogen.agentchat.contrib.retrieve_user_proxy_agent - INFO - Found 2 chunks.\u001b[0m\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "VectorDB returns doc_ids: [['bdfbc921', '7968cf3c']]\n", + "\u001b[32mAdding content of doc bdfbc921 to context.\u001b[0m\n", + "\u001b[32mAdding content of doc 7968cf3c to context.\u001b[0m\n", + "\u001b[33mragproxyagent\u001b[0m (to assistant):\n", + "\n", + "You're a retrieve augmented coding assistant. You answer user's questions based on your own knowledge and the\n", + "context provided by the user.\n", + "If you can't answer the question with or without the current context, you should reply exactly `UPDATE CONTEXT`.\n", + "For code generation, you must obey the following rules:\n", + "Rule 1. You MUST NOT install any packages because all the packages needed are already installed.\n", + "Rule 2. You must follow the formats below to write your code:\n", + "```language\n", + "# your code\n", + "```\n", + "\n", + "User's question is: How can I use FLAML to perform a classification task and use spark to do parallel training. Train 30 seconds and force cancel jobs if time limit is reached.\n", + "\n", + "Context is: # Integrate - Spark\n", + "\n", + "FLAML has integrated Spark for distributed training. There are two main aspects of integration with Spark:\n", + "\n", + "- Use Spark ML estimators for AutoML.\n", + "- Use Spark to run training in parallel spark jobs.\n", + "\n", + "## Spark ML Estimators\n", + "\n", + "FLAML integrates estimators based on Spark ML models. These models are trained in parallel using Spark, so we called them Spark estimators. To use these models, you first need to organize your data in the required format.\n", + "\n", + "### Data\n", + "\n", + "For Spark estimators, AutoML only consumes Spark data. FLAML provides a convenient function `to_pandas_on_spark` in the `flaml.automl.spark.utils` module to convert your data into a pandas-on-spark (`pyspark.pandas`) dataframe/series, which Spark estimators require.\n", + "\n", + "This utility function takes data in the form of a `pandas.Dataframe` or `pyspark.sql.Dataframe` and converts it into a pandas-on-spark dataframe. It also takes `pandas.Series` or `pyspark.sql.Dataframe` and converts it into a [pandas-on-spark](https://spark.apache.org/docs/latest/api/python/user_guide/pandas_on_spark/index.html) series. If you pass in a `pyspark.pandas.Dataframe`, it will not make any changes.\n", + "\n", + "This function also accepts optional arguments `index_col` and `default_index_type`.\n", + "\n", + "- `index_col` is the column name to use as the index, default is None.\n", + "- `default_index_type` is the default index type, default is \"distributed-sequence\". More info about default index type could be found on Spark official [documentation](https://spark.apache.org/docs/latest/api/python/user_guide/pandas_on_spark/options.html#default-index-type)\n", + "\n", + "Here is an example code snippet for Spark Data:\n", + "\n", + "```python\n", + "import pandas as pd\n", + "from flaml.automl.spark.utils import to_pandas_on_spark\n", + "\n", + "# Creating a dictionary\n", + "data = {\n", + " \"Square_Feet\": [800, 1200, 1800, 1500, 850],\n", + " \"Age_Years\": [20, 15, 10, 7, 25],\n", + " \"Price\": [100000, 200000, 300000, 240000, 120000],\n", + "}\n", + "\n", + "# Creating a pandas DataFrame\n", + "dataframe = pd.DataFrame(data)\n", + "label = \"Price\"\n", + "\n", + "# Convert to pandas-on-spark dataframe\n", + "psdf = to_pandas_on_spark(dataframe)\n", + "```\n", + "\n", + "To use Spark ML models you need to format your data appropriately. Specifically, use [`VectorAssembler`](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.feature.VectorAssembler.html) to merge all feature columns into a single vector column.\n", + "\n", + "Here is an example of how to use it:\n", + "\n", + "```python\n", + "from pyspark.ml.feature import VectorAssembler\n", + "\n", + "columns = psdf.columns\n", + "feature_cols = [col for col in columns if col != label]\n", + "featurizer = VectorAssembler(inputCols=feature_cols, outputCol=\"features\")\n", + "psdf = featurizer.transform(psdf.to_spark(index_col=\"index\"))[\"index\", \"features\"]\n", + "```\n", + "\n", + "Later in conducting the experiment, use your pandas-on-spark data like non-spark data and pass them using `X_train, y_train` or `dataframe, label`.\n", + "\n", + "### Estimators\n", + "\n", + "#### Model List\n", + "\n", + "- `lgbm_spark`: The class for fine-tuning Spark version LightGBM models, using [SynapseML](https://microsoft.github.io/SynapseML/docs/features/lightgbm/about/) API.\n", + "\n", + "#### Usage\n", + "\n", + "First, prepare your data in the required format as described in the previous section.\n", + "\n", + "By including the models you intend to try in the `estimators_list` argument to `flaml.automl`, FLAML will start trying configurations for these models. If your input is Spark data, FLAML will also use estimators with the `_spark` postfix by default, even if you haven't specified them.\n", + "\n", + "Here is an example code snippet using SparkML models in AutoML:\n", + "\n", + "```python\n", + "import flaml\n", + "\n", + "# prepare your data in pandas-on-spark format as we previously mentioned\n", + "\n", + "automl = flaml.AutoML()\n", + "settings = {\n", + " \"time_budget\": 30,\n", + " \"metric\": \"r2\",\n", + " \"estimator_list\": [\"lgbm_spark\"], # this setting is optional\n", + " \"task\": \"regression\",\n", + "}\n", + "\n", + "automl.fit(\n", + " dataframe=psdf,\n", + " label=label,\n", + " **settings,\n", + ")\n", + "```\n", + "\n", + "[Link to notebook](https://github.com/microsoft/FLAML/blob/main/notebook/automl_bankrupt_synapseml.ipynb) | [Open in colab](https://colab.research.google.com/github/microsoft/FLAML/blob/main/notebook/automl_bankrupt_synapseml.ipynb)\n", + "\n", + "## Parallel Spark Jobs\n", + "\n", + "You can activate Spark as the parallel backend during parallel tuning in both [AutoML](/docs/Use-Cases/Task-Oriented-AutoML#parallel-tuning) and [Hyperparameter Tuning](/docs/Use-Cases/Tune-User-Defined-Function#parallel-tuning), by setting the `use_spark` to `true`. FLAML will dispatch your job to the distributed Spark backend using [`joblib-spark`](https://github.com/joblib/joblib-spark).\n", + "\n", + "Please note that you should not set `use_spark` to `true` when applying AutoML and Tuning for Spark Data. This is because only SparkML models will be used for Spark Data in AutoML and Tuning. As SparkML models run in parallel, there is no need to distribute them with `use_spark` again.\n", + "\n", + "All the Spark-related arguments are stated below. These arguments are available in both Hyperparameter Tuning and AutoML:\n", + "\n", + "- `use_spark`: boolean, default=False | Whether to use spark to run the training in parallel spark jobs. This can be used to accelerate training on large models and large datasets, but will incur more overhead in time and thus slow down training in some cases. GPU training is not supported yet when use_spark is True. For Spark clusters, by default, we will launch one trial per executor. However, sometimes we want to launch more trials than the number of executors (e.g., local mode). In this case, we can set the environment variable `FLAML_MAX_CONCURRENT` to override the detected `num_executors`. The final number of concurrent trials will be the minimum of `n_concurrent_trials` and `num_executors`.\n", + "- `n_concurrent_trials`: int, default=1 | The number of concurrent trials. When n_concurrent_trials > 1, FLAML performes parallel tuning.\n", + "- `force_cancel`: boolean, default=False | Whether to forcely cancel Spark jobs if the search time exceeded the time budget. Spark jobs include parallel tuning jobs and Spark-based model training jobs.\n", + "\n", + "An example code snippet for using parallel Spark jobs:\n", + "\n", + "```python\n", + "import flaml\n", + "\n", + "automl_experiment = flaml.AutoML()\n", + "automl_settings = {\n", + " \"time_budget\": 30,\n", + " \"metric\": \"r2\",\n", + " \"task\": \"regression\",\n", + " \"n_concurrent_trials\": 2,\n", + " \"use_spark\": True,\n", + " \"force_cancel\": True, # Activating the force_cancel option can immediately halt Spark jobs once they exceed the allocated time_budget.\n", + "}\n", + "\n", + "automl.fit(\n", + " dataframe=dataframe,\n", + " label=label,\n", + " **automl_settings,\n", + ")\n", + "```\n", + "\n", + "[Link to notebook](https://github.com/microsoft/FLAML/blob/main/notebook/integrate_spark.ipynb) | [Open in colab](https://colab.research.google.com/github/microsoft/FLAML/blob/main/notebook/integrate_spark.ipynb)\n", + "# Research\n", + "\n", + "For technical details, please check our research publications.\n", + "\n", + "- [FLAML: A Fast and Lightweight AutoML Library](https://www.microsoft.com/en-us/research/publication/flaml-a-fast-and-lightweight-automl-library/). Chi Wang, Qingyun Wu, Markus Weimer, Erkang Zhu. MLSys 2021.\n", + "\n", + "```bibtex\n", + "@inproceedings{wang2021flaml,\n", + " title={FLAML: A Fast and Lightweight AutoML Library},\n", + " author={Chi Wang and Qingyun Wu and Markus Weimer and Erkang Zhu},\n", + " year={2021},\n", + " booktitle={MLSys},\n", + "}\n", + "```\n", + "\n", + "- [Frugal Optimization for Cost-related Hyperparameters](https://arxiv.org/abs/2005.01571). Qingyun Wu, Chi Wang, Silu Huang. AAAI 2021.\n", + "\n", + "```bibtex\n", + "@inproceedings{wu2021cfo,\n", + " title={Frugal Optimization for Cost-related Hyperparameters},\n", + " author={Qingyun Wu and Chi Wang and Silu Huang},\n", + " year={2021},\n", + " booktitle={AAAI},\n", + "}\n", + "```\n", + "\n", + "- [Economical Hyperparameter Optimization With Blended Search Strategy](https://www.microsoft.com/en-us/research/publication/economical-hyperparameter-optimization-with-blended-search-strategy/). Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. ICLR 2021.\n", + "\n", + "```bibtex\n", + "@inproceedings{wang2021blendsearch,\n", + " title={Economical Hyperparameter Optimization With Blended Search Strategy},\n", + " author={Chi Wang and Qingyun Wu and Silu Huang and Amin Saied},\n", + " year={2021},\n", + " booktitle={ICLR},\n", + "}\n", + "```\n", + "\n", + "- [An Empirical Study on Hyperparameter Optimization for Fine-Tuning Pre-trained Language Models](https://aclanthology.org/2021.acl-long.178.pdf). Susan Xueqing Liu, Chi Wang. ACL 2021.\n", + "\n", + "```bibtex\n", + "@inproceedings{liuwang2021hpolm,\n", + " title={An Empirical Study on Hyperparameter Optimization for Fine-Tuning Pre-trained Language Models},\n", + " author={Susan Xueqing Liu and Chi Wang},\n", + " year={2021},\n", + " booktitle={ACL},\n", + "}\n", + "```\n", + "\n", + "- [ChaCha for Online AutoML](https://www.microsoft.com/en-us/research/publication/chacha-for-online-automl/). Qingyun Wu, Chi Wang, John Langford, Paul Mineiro and Marco Rossi. ICML 2021.\n", + "\n", + "```bibtex\n", + "@inproceedings{wu2021chacha,\n", + " title={ChaCha for Online AutoML},\n", + " author={Qingyun Wu and Chi Wang and John Langford and Paul Mineiro and Marco Rossi},\n", + " year={2021},\n", + " booktitle={ICML},\n", + "}\n", + "```\n", + "\n", + "- [Fair AutoML](https://arxiv.org/abs/2111.06495). Qingyun Wu, Chi Wang. ArXiv preprint arXiv:2111.06495 (2021).\n", + "\n", + "```bibtex\n", + "@inproceedings{wuwang2021fairautoml,\n", + " title={Fair AutoML},\n", + " author={Qingyun Wu and Chi Wang},\n", + " year={2021},\n", + " booktitle={ArXiv preprint arXiv:2111.06495},\n", + "}\n", + "```\n", + "\n", + "- [Mining Robust Default Configurations for Resource-constrained AutoML](https://arxiv.org/abs/2202.09927). Moe Kayali, Chi Wang. ArXiv preprint arXiv:2202.09927 (2022).\n", + "\n", + "```bibtex\n", + "@inproceedings{kayaliwang2022default,\n", + " title={Mining Robust Default Configurations for Resource-constrained AutoML},\n", + " author={Moe Kayali and Chi Wang},\n", + " year={2022},\n", + " booktitle={ArXiv preprint arXiv:2202.09927},\n", + "}\n", + "```\n", + "\n", + "- [Targeted Hyperparameter Optimization with Lexicographic Preferences Over Multiple Objectives](https://openreview.net/forum?id=0Ij9_q567Ma). Shaokun Zhang, Feiran Jia, Chi Wang, Qingyun Wu. ICLR 2023 (notable-top-5%).\n", + "\n", + "```bibtex\n", + "@inproceedings{zhang2023targeted,\n", + " title={Targeted Hyperparameter Optimization with Lexicographic Preferences Over Multiple Objectives},\n", + " author={Shaokun Zhang and Feiran Jia and Chi Wang and Qingyun Wu},\n", + " booktitle={International Conference on Learning Representations},\n", + " year={2023},\n", + " url={https://openreview.net/forum?id=0Ij9_q567Ma},\n", + "}\n", + "```\n", + "\n", + "- [Cost-Effective Hyperparameter Optimization for Large Language Model Generation Inference](https://arxiv.org/abs/2303.04673). Chi Wang, Susan Xueqing Liu, Ahmed H. Awadallah. ArXiv preprint arXiv:2303.04673 (2023).\n", + "\n", + "```bibtex\n", + "@inproceedings{wang2023EcoOptiGen,\n", + " title={Cost-Effective Hyperparameter Optimization for Large Language Model Generation Inference},\n", + " author={Chi Wang and Susan Xueqing Liu and Ahmed H. Awadallah},\n", + " year={2023},\n", + " booktitle={ArXiv preprint arXiv:2303.04673},\n", + "}\n", + "```\n", + "\n", + "- [An Empirical Study on Challenging Math Problem Solving with GPT-4](https://arxiv.org/abs/2306.01337). Yiran Wu, Feiran Jia, Shaokun Zhang, Hangyu Li, Erkang Zhu, Yue Wang, Yin Tat Lee, Richard Peng, Qingyun Wu, Chi Wang. ArXiv preprint arXiv:2306.01337 (2023).\n", + "\n", + "```bibtex\n", + "@inproceedings{wu2023empirical,\n", + " title={An Empirical Study on Challenging Math Problem Solving with GPT-4},\n", + " author={Yiran Wu and Feiran Jia and Shaokun Zhang and Hangyu Li and Erkang Zhu and Yue Wang and Yin Tat Lee and Richard Peng and Qingyun Wu and Chi Wang},\n", + " year={2023},\n", + " booktitle={ArXiv preprint arXiv:2306.01337},\n", + "}\n", + "```\n", + "\n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33massistant\u001b[0m (to ragproxyagent):\n", + "\n", + "```python\n", + "import pandas as pd\n", + "from pyspark.ml.feature import VectorAssembler\n", + "import flaml\n", + "from flaml.automl.spark.utils import to_pandas_on_spark\n", + "\n", + "# Creating a dictionary for the example data\n", + "data = {\n", + " \"Square_Feet\": [800, 1200, 1800, 1500, 850],\n", + " \"Age_Years\": [20, 15, 10, 7, 25],\n", + " \"Price\": [100000, 200000, 300000, 240000, 120000],\n", + "}\n", + "\n", + "# Creating a pandas DataFrame\n", + "dataframe = pd.DataFrame(data)\n", + "label = \"Price\"\n", + "\n", + "# Convert to pandas-on-spark dataframe\n", + "psdf = to_pandas_on_spark(dataframe)\n", + "\n", + "# Prepare features using VectorAssembler\n", + "columns = psdf.columns\n", + "feature_cols = [col for col in columns if col != label]\n", + "featurizer = VectorAssembler(inputCols=feature_cols, outputCol=\"features\")\n", + "psdf = featurizer.transform(psdf.to_spark(index_col=\"index\"))[[\"index\", \"features\"]]\n", + "\n", + "# Setting up and running FLAML for AutoML with Spark\n", + "automl = flaml.AutoML()\n", + "automl_settings = {\n", + " \"time_budget\": 30, # Set the time budget to 30 seconds\n", + " \"metric\": \"r2\", # Performance metric\n", + " \"task\": \"regression\", # Problem type\n", + " \"n_concurrent_trials\": 2, # Number of concurrent trials\n", + " \"use_spark\": True, # Use Spark for parallel jobs\n", + " \"force_cancel\": True, # Force cancel jobs if time limit is reached\n", + "}\n", + "\n", + "automl.fit(\n", + " dataframe=psdf,\n", + " label=label,\n", + " **automl_settings\n", + ")\n", + "```\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mragproxyagent\u001b[0m (to assistant):\n", + "\n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33massistant\u001b[0m (to ragproxyagent):\n", + "\n", + "UPDATE CONTEXT\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[32mUpdating context and resetting conversation.\u001b[0m\n", + "VectorDB returns doc_ids: [['bdfbc921', '7968cf3c']]\n", + "\u001b[32mNo more context, will terminate.\u001b[0m\n", + "\u001b[33mragproxyagent\u001b[0m (to assistant):\n", + "\n", + "TERMINATE\n", + "\n", + "--------------------------------------------------------------------------------\n" + ] + } + ], + "source": [ + "# reset the assistant. Always reset the assistant before starting a new conversation.\n", + "assistant.reset()\n", + "\n", + "# given a problem, we use the ragproxyagent to generate a prompt to be sent to the assistant as the initial message.\n", + "# the assistant receives the message and generates a response. The response will be sent back to the ragproxyagent for processing.\n", + "# The conversation continues until the termination condition is met, in RetrieveChat, the termination condition when no human-in-loop is no code block detected.\n", + "# With human-in-loop, the conversation will continue until the user says \"exit\".\n", + "code_problem = \"How can I use FLAML to perform a classification task and use spark to do parallel training. Train 30 seconds and force cancel jobs if time limit is reached.\"\n", + "chat_result = ragproxyagent.initiate_chat(assistant, message=ragproxyagent.message_generator, problem=code_problem)" + ] + } + ], + "metadata": { + "front_matter": { + "description": "Explore the use of AutoGen's RetrieveChat for tasks like code generation from docstrings, answering complex questions with human feedback, and exploiting features like Update Context, custom prompts, and few-shot learning.", + "tags": [ + "RAG" + ] + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.11.7" + }, + "skip_test": "Requires interactive usage" + }, + "nbformat": 4, + "nbformat_minor": 4 + } \ No newline at end of file diff --git a/setup.py b/setup.py index a5a71db02a..92b6b29eb5 100644 --- a/setup.py +++ b/setup.py @@ -93,6 +93,7 @@ "retrievechat-pgvector": retrieve_chat_pgvector, "retrievechat-mongodb": [*retrieve_chat, "pymongo>=4.0.0"], "retrievechat-qdrant": [*retrieve_chat, "qdrant_client", "fastembed>=0.3.1"], + "retrievechat-couchbase": [*retrieve_chat, "couchbase>=4.3.0"], "graph_rag_falkor_db": graph_rag_falkor_db, "autobuild": autobuild, "captainagent": autobuild + ["pandas"], diff --git a/test/agentchat/contrib/vectordb/test_couchbase.py b/test/agentchat/contrib/vectordb/test_couchbase.py new file mode 100644 index 0000000000..23279d8d89 --- /dev/null +++ b/test/agentchat/contrib/vectordb/test_couchbase.py @@ -0,0 +1,153 @@ +import logging +import os +import random +from time import sleep + +import pytest +from dotenv import load_dotenv + +try: + + import couchbase + import sentence_transformers + import sys + import os + + from autogen.agentchat.contrib.vectordb.couchbase import CouchbaseVectorDB +except ImportError: + print("skipping test_couchbase.py. It requires one to pip install couchbase or the extra [retrievechat-couchbase]") + logger = logging.getLogger(__name__) + logger.warning( + f"skipping {__name__}. It requires one to pip install couchbase or the extra [retrievechat-couchbase]" + ) + pytest.skip("Required modules not installed", allow_module_level=True) + +from couchbase.auth import PasswordAuthenticator +from couchbase.cluster import Cluster, ClusterOptions + +logger = logging.getLogger(__name__) + +# Get the directory of the current script +script_dir = os.path.dirname(os.path.abspath(__file__)) + +# Construct the absolute path to the .env file +env_path = os.path.join(script_dir, ".env") + +# Load the .env file +load_dotenv(env_path) +load_dotenv(".env") + +COUCHBASE_HOST = os.environ.get("CB_CONN_STR", "couchbase://localhost") +COUCHBASE_USERNAME = os.environ.get("CB_USERNAME", "Administrator") +COUCHBASE_PASSWORD = os.environ.get("CB_PASSWORD", "password") +COUCHBASE_BUCKET = os.environ.get("CB_BUCKET", "autogen_test_bucket") +COUCHBASE_SCOPE = os.environ.get("CB_SCOPE", "_default") +COUCHBASE_COLLECTION = os.environ.get("CB_COLLECTION", "autogen_test_vectorstore") +COUCHBASE_INDEX = os.environ.get("CB_INDEX_NAME", "vector_index") + +RETRIES = 10 +DELAY = 2 +TIMEOUT = 120.0 + + +def _empty_collections_and_delete_indexes(cluster: Cluster, bucket_name, scope_name, collections=None): + bucket = cluster.bucket(bucket_name) + try: + scope_manager = bucket.collections().get_all_scopes(scope_name=scope_name) + for scope_ in scope_manager: + all_collections = scope_.collections + for curr_collection in all_collections: + bucket.collections().drop_collection(scope_name, curr_collection.name) + except Exception as e: + logger.warning(f"Failed to drop collections: {e}") + + +@pytest.fixture +def db(): + print("Creating couchbase connection", COUCHBASE_HOST, COUCHBASE_USERNAME, COUCHBASE_PASSWORD) + cluster = Cluster(COUCHBASE_HOST, ClusterOptions(PasswordAuthenticator(COUCHBASE_USERNAME, COUCHBASE_PASSWORD))) + _empty_collections_and_delete_indexes(cluster, COUCHBASE_BUCKET, COUCHBASE_SCOPE) + vectorstore = CouchbaseVectorDB( + connection_string=COUCHBASE_HOST, + username=COUCHBASE_USERNAME, + password=COUCHBASE_PASSWORD, + bucket_name=COUCHBASE_BUCKET, + scope_name=COUCHBASE_SCOPE, + collection_name=COUCHBASE_COLLECTION, + index_name=COUCHBASE_INDEX, + ) + yield vectorstore + _empty_collections_and_delete_indexes(cluster, COUCHBASE_BUCKET, COUCHBASE_SCOPE) + + +_COLLECTION_NAMING_CACHE = [] + + +@pytest.fixture +def collection_name(): + collection_id = random.randint(0, 100) + while collection_id in _COLLECTION_NAMING_CACHE: + collection_id = random.randint(0, 100) + _COLLECTION_NAMING_CACHE.append(collection_id) + return f"{COUCHBASE_COLLECTION}_{collection_id}" + + +def test_couchbase(db, collection_name): + # db = CouchbaseVectorDB(path=".db") + with pytest.raises(Exception): + curr_col = db.get_collection(collection_name) + curr_col.upsert("1", {"content": "Dogs are lovely."}) + + collection = db.create_collection(collection_name, overwrite=True, get_or_create=True) + assert collection.name == collection_name + collection.upsert("1", {"content": "Dogs are lovely."}) + + # test_delete_collection + db.delete_collection(collection_name) + sleep(5) # wait for the collection to be deleted + with pytest.raises(Exception): + curr_col = db.get_collection(collection_name) + curr_col.upsert("1", {"content": "Dogs are lovely."}) + + # test more create collection + collection = db.create_collection(collection_name, overwrite=False, get_or_create=False) + assert collection.name == collection_name + pytest.raises(ValueError, db.create_collection, collection_name, overwrite=False, get_or_create=False) + collection = db.create_collection(collection_name, overwrite=True, get_or_create=False) + assert collection.name == collection_name + collection = db.create_collection(collection_name, overwrite=False, get_or_create=True) + assert collection.name == collection_name + + # test_get_collection + collection = db.get_collection(collection_name) + assert collection.name == collection_name + + # test_insert_docs + docs = [{"content": "doc1", "id": "1"}, {"content": "doc2", "id": "2"}, {"content": "doc3", "id": "3"}] + db.insert_docs(docs, collection_name, upsert=False) + res = db.get_collection(collection_name).get_multi(["1", "2"]).results + + assert res["1"].value["content"] == "doc1" + assert res["2"].value["content"] == "doc2" + + # test_update_docs + docs = [{"content": "doc11", "id": "1"}, {"content": "doc2", "id": "2"}, {"content": "doc3", "id": "3"}] + db.update_docs(docs, collection_name) + res = db.get_collection(collection_name).get_multi(["1", "2"]).results + assert res["1"].value["content"] == "doc11" + assert res["2"].value["content"] == "doc2" + + # test_delete_docs + ids = ["1"] + db.delete_docs(ids, collection_name) + with pytest.raises(Exception): + res = db.get_collection(collection_name).get(ids[0]) + + # test_retrieve_docs + queries = ["doc2", "doc3"] + res = db.retrieve_docs(queries, collection_name) + texts = [[item[0]["content"] for item in sublist] for sublist in res] + received_ids = [[item[0]["id"] for item in sublist] for sublist in res] + + assert texts[0] == ["doc2", "doc3"] + assert received_ids[0] == ["2", "3"] \ No newline at end of file diff --git a/website/blog/2023-10-18-RetrieveChat/index.mdx b/website/blog/2023-10-18-RetrieveChat/index.mdx index abc95a3fe1..a82799484c 100644 --- a/website/blog/2023-10-18-RetrieveChat/index.mdx +++ b/website/blog/2023-10-18-RetrieveChat/index.mdx @@ -200,8 +200,8 @@ ragproxyagent = RetrieveUserProxyAgent( ### Customizing Vector Database -We are using chromadb as the default vector database, you can also use mongodb, pgvectordb and qdrantdb -by simply set `vector_db` to `mongodb`, `pgvector` and `qdrant` in `retrieve_config`, respectively. +We are using chromadb as the default vector database, you can also use mongodb, pgvectordb, qdrantdb and couchbase +by simply set `vector_db` to `mongodb`, `pgvector`, `qdrant` and `couchbase` in `retrieve_config`, respectively. To plugin any other dbs, you can also extend class `agentchat.contrib.vectordb.base`, check out the code [here](https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/contrib/vectordb/base.py). @@ -404,3 +404,4 @@ You can check out more example notebooks for RAG use cases: - [Using RetrieveChat with Qdrant for Retrieve Augmented Code Generation and Question Answering](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_qdrant.ipynb) - [Using RetrieveChat Powered by PGVector for Retrieve Augmented Code Generation and Question Answering](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_pgvector.ipynb) - [Using RetrieveChat Powered by MongoDB Atlas for Retrieve Augmented Code Generation and Question Answering](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_mongodb.ipynb) +- [Using RetrieveChat Powered by Couchbase for Retrieve Augmented Code Generation and Question Answering](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_couchbase.ipynb) \ No newline at end of file diff --git a/website/docs/topics/retrieval_augmentation.md b/website/docs/topics/retrieval_augmentation.md index 22ee7d28ba..033db3873d 100644 --- a/website/docs/topics/retrieval_augmentation.md +++ b/website/docs/topics/retrieval_augmentation.md @@ -126,6 +126,7 @@ For more detailed examples and notebooks showcasing the usage of retrieval augme - Automated Code Generation and Question Answering with [PGVector](https://github.com/pgvector/pgvector) based Retrieval Augmented Agents - [View Notebook](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_pgvector.ipynb) - Automated Code Generation and Question Answering with [Qdrant](https://qdrant.tech/) based Retrieval Augmented Agents - [View Notebook](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_qdrant.ipynb) - Automated Code Generation and Question Answering with [MongoDB Atlas](https://www.mongodb.com/) based Retrieval Augmented Agents - [View Notebook](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_mongodb.ipynb) +- Automated Code Generation and Question Answering with [Couchbase](https://www.couchbase.com/) based Retrieval Augmented Agents - [View Notebook](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_RetrieveChat_couchbase.ipynb) - Chat with OpenAI Assistant with Retrieval Augmentation - [View Notebook](https://github.com/ag2ai/ag2/blob/main/notebook/agentchat_oai_assistant_retrieval.ipynb) - **RAG**: Group Chat with Retrieval Augmented Generation (with 5 group member agents and 1 manager agent) - [View Notebook](/docs/notebooks/agentchat_groupchat_RAG)