归档

graphrag/config/enums.py

@@ -98,6 +98,7 @@ class LLMType(str, Enum):
     # Embeddings
     OpenAIEmbedding = "openai_embedding"
     AzureOpenAIEmbedding = "azure_openai_embedding"
+    DashScopeEmbedding = "dashscope_embedding"
 
     # Raw Completion
     OpenAI = "openai"

graphrag/index/llm/load_llm.py

@@ -16,6 +16,7 @@
     LLMCache,
     LLMLimiter,
     MockCompletionLLM,
+    DashScopeEmbeddingLLM,
     OpenAIConfiguration,
     create_openai_chat_llm,
     create_openai_client,
@@ -94,6 +95,13 @@ def on_error(
 
     return on_error
 
+def _load_dashscope_embeddings_llm(
+    on_error: ErrorHandlerFn,
+    cache: LLMCache,
+    config: dict[str, Any],
+    azure=False,
+): 
+    return DashScopeEmbeddingLLM(config)
 
 def _load_openai_completion_llm(
     on_error: ErrorHandlerFn,
@@ -227,6 +235,10 @@ def _load_static_response(
         "load": _load_azure_openai_chat_llm,
         "chat": True,
     },
+    LLMType.DashScopeEmbedding: {
+        "load": _load_dashscope_embeddings_llm,
+        "chat": False,
+    },
     LLMType.OpenAIEmbedding: {
         "load": _load_openai_embeddings_llm,
         "chat": False,

graphrag/llm/__init__.py

@@ -42,6 +42,7 @@
     LLMOutput,
     OnCacheActionFn,
 )
+from .dashscope import DashScopeEmbeddingLLM
 
 __all__ = [
     # LLM Types
@@ -88,4 +89,6 @@
     "create_openai_embedding_llm",
     # Limiters
     "create_tpm_rpm_limiters",
+    # DashScope
+    "DashScopeEmbeddingLLM"
 ]

graphrag/llm/dashscope/__init__.py

@@ -0,0 +1,5 @@
+from .dashscope_embedding import DashScopeEmbeddingLLM
+
+__all__ = [
+    "DashScopeEmbeddingLLM",
+]

graphrag/llm/dashscope/dashscope_embedding.py

@@ -0,0 +1,35 @@
+"""The EmbeddingsLLM class."""
+
+import logging
+import dashscope
+from typing_extensions import Unpack
+
+from graphrag.llm.base import BaseLLM
+from graphrag.llm.types import (
+    EmbeddingInput,
+    EmbeddingOutput,
+    LLMInput,
+)
+
+log = logging.getLogger(__name__)
+
+class DashScopeEmbeddingLLM(BaseLLM[EmbeddingInput, EmbeddingOutput]):
+    """A text-embedding generator LLM."""
+    def __init__(self, llm_config: dict):
+        log.info(f"llm_config: {llm_config}")
+        self.llm_config = llm_config or {}
+        self.api_key = self.llm_config.get("api_key", "")
+        self.model = self.llm_config.get("model", "text-embedding-v2")
+
+    async def _execute_llm(
+            self, input: EmbeddingInput, **kwargs: Unpack[LLMInput]
+    ) -> EmbeddingOutput:
+        log.info(f"input: {input}")
+
+        response = dashscope.TextEmbedding.call(
+            model=self.model,
+            input=input,
+            api_key=self.api_key
+        )
+
+        return [embedding["embedding"] for embedding in response.output["embeddings"]]

pyproject.toml

@@ -22,6 +22,11 @@ license = "MIT"
 readme = "README.md"
 packages = [{ include = "graphrag" }]
 
+[[tool.poetry.source]]
+name = "tsinghua"
+url = "https://pypi.tuna.tsinghua.edu.cn/simple"
+priority = "primary"
+
 [tool.poetry-dynamic-versioning]
 enable = true
 style = "pep440"
@@ -86,6 +91,7 @@ typing-extensions = "^4.12.2"
 azure-storage-blob = "^12.19.0"
 azure-identity = "^1.17.1"
 json-repair = "^0.25.3"
+dashscope = "^1.20.3"
 
 [tool.poetry.group.dev.dependencies]
 coverage = "^7.6.0"