begin rag project with llama index

request_llms/embed_models/bridge_all_embed.py

@@ -0,0 +1,40 @@
+import tiktoken, copy, re
+from functools import lru_cache
+from concurrent.futures import ThreadPoolExecutor
+from toolbox import get_conf, trimmed_format_exc, apply_gpt_academic_string_mask, read_one_api_model_name
+
+# Endpoint 重定向
+API_URL_REDIRECT, AZURE_ENDPOINT, AZURE_ENGINE = get_conf("API_URL_REDIRECT", "AZURE_ENDPOINT", "AZURE_ENGINE")
+openai_endpoint = "https://api.openai.com/v1/chat/completions"
+if not AZURE_ENDPOINT.endswith('/'): AZURE_ENDPOINT += '/'
+azure_endpoint = AZURE_ENDPOINT + f'openai/deployments/{AZURE_ENGINE}/chat/completions?api-version=2023-05-15'
+
+
+if openai_endpoint in API_URL_REDIRECT: openai_endpoint = API_URL_REDIRECT[openai_endpoint]
+
+openai_embed_endpoint = openai_endpoint.replace("chat/completions", "embeddings")
+
+from .openai_embed import OpenAiEmbeddingModel
+
+embed_model_info = {
+    # text-embedding-3-small    Increased performance over 2nd generation ada embedding model  |  1,536
+    "text-embedding-3-small": {
+        "embed_class": OpenAiEmbeddingModel,
+        "embed_endpoint": openai_embed_endpoint,
+        "embed_dimension": 1536,
+    },
+
+    # text-embedding-3-large    Most capable embedding model for both english and non-english tasks  |   3,072
+    "text-embedding-3-large": {
+        "embed_class": OpenAiEmbeddingModel,
+        "embed_endpoint": openai_embed_endpoint,
+        "embed_dimension": 3072,
+    },
+
+    # text-embedding-ada-002    Most capable 2nd generation embedding model, replacing 16 first generation models   |  1,536
+    "text-embedding-ada-002": {
+        "embed_class": OpenAiEmbeddingModel,
+        "embed_endpoint": openai_embed_endpoint,
+        "embed_dimension": 1536,
+    },
+}

request_llms/embed_models/openai_embed.py

@@ -0,0 +1,60 @@
+from llama_index.embeddings.openai import OpenAIEmbedding
+from openai import OpenAI
+from toolbox import get_conf
+from toolbox import CatchException, update_ui, get_conf, select_api_key, get_log_folder, ProxyNetworkActivate
+from shared_utils.key_pattern_manager import select_api_key_for_embed_models
+from typing import List, Any
+
+class OpenAiEmbeddingModel():
+
+    def __init__(self, llm_kwargs:dict=None):
+        self.llm_kwargs = llm_kwargs
+
+    def compute_embedding(self, text="这是要计算嵌入的文本", llm_kwargs:dict=None, batch_mode=False):
+        from .bridge_all_embed import embed_model_info
+
+        # load kwargs
+        if llm_kwargs is None:
+            llm_kwargs = self.llm_kwargs
+        if llm_kwargs is None:
+            raise RuntimeError("llm_kwargs is not provided!")
+
+        # setup api and req url
+        api_key = select_api_key_for_embed_models(llm_kwargs['api_key'], llm_kwargs['llm_model'])
+        embed_model = llm_kwargs['llm_model']
+        base_url = embed_model_info[llm_kwargs['llm_model']]['embed_endpoint'].replace('embeddings', '')
+
+        # send and compute
+        with ProxyNetworkActivate("Connect_OpenAI_Embedding"):
+            self.oai_client = OpenAI(api_key=api_key, base_url=base_url)
+            if batch_mode:
+                input = text
+                assert isinstance(text, list)
+            else:
+                input = [text]
+                assert isinstance(text, str)
+            res = self.oai_client.embeddings.create(input=input, model=embed_model)
+
+        # parse result
+        if batch_mode:
+            embedding = [d.embedding for d in res.data]
+        else:
+            embedding = res.data[0].embedding
+
+        return embedding
+
+
+    def embedding_dimension(self, llm_kwargs):
+
+        from .bridge_all_embed import embed_model_info
+        return embed_model_info[llm_kwargs['llm_model']]['embed_dimension']
+    
+    def get_text_embedding_batch(
+            self,
+            texts: List[str],
+            show_progress: bool = False,
+        ):
+            return self.compute_embedding(texts, batch_mode=True)
+
+if __name__ == "__main__":
+    pass