Merge branch 'main' into feat_werewolf

metagpt/rag/factories/base.py

@@ -26,6 +26,9 @@ class GenericFactory:
         if creator:
             return creator(**kwargs)
 
+        self._raise_for_key(key)
+
+    def _raise_for_key(self, key: Any):
         raise ValueError(f"Creator not registered for key: {key}")
 
 

metagpt/rag/factories/embedding.py

@@ -1,37 +1,103 @@
 """RAG Embedding Factory."""
+from __future__ import annotations
+
+from typing import Any
 
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
+from llama_index.embeddings.gemini import GeminiEmbedding
+from llama_index.embeddings.ollama import OllamaEmbedding
 from llama_index.embeddings.openai import OpenAIEmbedding
 
 from metagpt.config2 import config
+from metagpt.configs.embedding_config import EmbeddingType
 from metagpt.configs.llm_config import LLMType
 from metagpt.rag.factories.base import GenericFactory
 
 
 class RAGEmbeddingFactory(GenericFactory):
-    """Create LlamaIndex Embedding with MetaGPT's config."""
+    """Create LlamaIndex Embedding with MetaGPT's embedding config."""
 
     def __init__(self):
         creators = {
+            EmbeddingType.OPENAI: self._create_openai,
+            EmbeddingType.AZURE: self._create_azure,
+            EmbeddingType.GEMINI: self._create_gemini,
+            EmbeddingType.OLLAMA: self._create_ollama,
+            # For backward compatibility
             LLMType.OPENAI: self._create_openai,
             LLMType.AZURE: self._create_azure,
         }
         super().__init__(creators)
 
-    def get_rag_embedding(self, key: LLMType = None) -> BaseEmbedding:
-        """Key is LLMType, default use config.llm.api_type."""
-        return super().get_instance(key or config.llm.api_type)
+    def get_rag_embedding(self, key: EmbeddingType = None) -> BaseEmbedding:
+        """Key is EmbeddingType."""
+        return super().get_instance(key or self._resolve_embedding_type())
 
-    def _create_openai(self):
-        return OpenAIEmbedding(api_key=config.llm.api_key, api_base=config.llm.base_url)
+    def _resolve_embedding_type(self) -> EmbeddingType | LLMType:
+        """Resolves the embedding type.
 
-    def _create_azure(self):
-        return AzureOpenAIEmbedding(
-            azure_endpoint=config.llm.base_url,
-            api_key=config.llm.api_key,
-            api_version=config.llm.api_version,
+        If the embedding type is not specified, for backward compatibility, it checks if the LLM API type is either OPENAI or AZURE.
+        Raise TypeError if embedding type not found.
+        """
+        if config.embedding.api_type:
+            return config.embedding.api_type
+
+        if config.llm.api_type in [LLMType.OPENAI, LLMType.AZURE]:
+            return config.llm.api_type
+
+        raise TypeError("To use RAG, please set your embedding in config2.yaml.")
+
+    def _create_openai(self) -> OpenAIEmbedding:
+        params = dict(
+            api_key=config.embedding.api_key or config.llm.api_key,
+            api_base=config.embedding.base_url or config.llm.base_url,
         )
 
+        self._try_set_model_and_batch_size(params)
+
+        return OpenAIEmbedding(**params)
+
+    def _create_azure(self) -> AzureOpenAIEmbedding:
+        params = dict(
+            api_key=config.embedding.api_key or config.llm.api_key,
+            azure_endpoint=config.embedding.base_url or config.llm.base_url,
+            api_version=config.embedding.api_version or config.llm.api_version,
+        )
+
+        self._try_set_model_and_batch_size(params)
+
+        return AzureOpenAIEmbedding(**params)
+
+    def _create_gemini(self) -> GeminiEmbedding:
+        params = dict(
+            api_key=config.embedding.api_key,
+            api_base=config.embedding.base_url,
+        )
+
+        self._try_set_model_and_batch_size(params)
+
+        return GeminiEmbedding(**params)
+
+    def _create_ollama(self) -> OllamaEmbedding:
+        params = dict(
+            base_url=config.embedding.base_url,
+        )
+
+        self._try_set_model_and_batch_size(params)
+
+        return OllamaEmbedding(**params)
+
+    def _try_set_model_and_batch_size(self, params: dict):
+        """Set the model_name and embed_batch_size only when they are specified."""
+        if config.embedding.model:
+            params["model_name"] = config.embedding.model
+
+        if config.embedding.embed_batch_size:
+            params["embed_batch_size"] = config.embedding.embed_batch_size
+
+    def _raise_for_key(self, key: Any):
+        raise ValueError(f"The embedding type is currently not supported: `{type(key)}`, {key}")
+
 
 get_rag_embedding = RAGEmbeddingFactory().get_rag_embedding

metagpt/rag/factories/llm.py

@@ -1,5 +1,5 @@
 """RAG LLM."""
-
+import asyncio
 from typing import Any
 
 from llama_index.core.constants import DEFAULT_CONTEXT_WINDOW
@@ -15,7 +15,7 @@ from pydantic import Field
 from metagpt.config2 import config
 from metagpt.llm import LLM
 from metagpt.provider.base_llm import BaseLLM
-from metagpt.utils.async_helper import run_coroutine_in_new_loop
+from metagpt.utils.async_helper import NestAsyncio
 from metagpt.utils.token_counter import TOKEN_MAX
 
 
@@ -39,7 +39,8 @@ class RAGLLM(CustomLLM):
 
     @llm_completion_callback()
     def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        return run_coroutine_in_new_loop(self.acomplete(prompt, **kwargs))
+        NestAsyncio.apply_once()
+        return asyncio.get_event_loop().run_until_complete(self.acomplete(prompt, **kwargs))
 
     @llm_completion_callback()
     async def acomplete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:

metagpt/rag/schema.py

@@ -1,15 +1,17 @@
 """RAG schemas."""
 
 from pathlib import Path
-from typing import Any, Literal, Optional, Union
+from typing import Any, ClassVar, Literal, Optional, Union
 
 from chromadb.api.types import CollectionMetadata
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.core.indices.base import BaseIndex
 from llama_index.core.schema import TextNode
 from llama_index.core.vector_stores.types import VectorStoreQueryMode
-from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
 
+from metagpt.config2 import config
+from metagpt.configs.embedding_config import EmbeddingType
 from metagpt.rag.interface import RAGObject
 
 
@@ -32,7 +34,19 @@ class IndexRetrieverConfig(BaseRetrieverConfig):
 class FAISSRetrieverConfig(IndexRetrieverConfig):
     """Config for FAISS-based retrievers."""
 
-    dimensions: int = Field(default=1536, description="Dimensionality of the vectors for FAISS index construction.")
+    dimensions: int = Field(default=0, description="Dimensionality of the vectors for FAISS index construction.")
+
+    _embedding_type_to_dimensions: ClassVar[dict[EmbeddingType, int]] = {
+        EmbeddingType.GEMINI: 768,
+        EmbeddingType.OLLAMA: 4096,
+    }
+
+    @model_validator(mode="after")
+    def check_dimensions(self):
+        if self.dimensions == 0:
+            self.dimensions = self._embedding_type_to_dimensions.get(config.embedding.api_type, 1536)
+
+        return self
 
 
 class BM25RetrieverConfig(IndexRetrieverConfig):