Add Local-LLM embedder (#1032)

* Add Local-LLM embedder * Add error handling
Josh-XT · Oct 4, 2023 · 9df047d · 9df047d
1 parent 8220a9e
commit 9df047d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 13 deletions.
diff --git a/agixt/Embedding.py b/agixt/Embedding.py
@@ -25,15 +25,25 @@ def __init__(self, agent_settings=None):
             or self.agent_settings["embedder"] == "default"
         ):
             self.agent_settings["embedder"] = "default"
-            self.embedder = self.default_embedder
-            self.chunk_size = 256
-        else:
+            if "provider" in self.agent_settings:
+                if "provider" == "local":
+                    self.agent_settings["embedder"] = "local"
+                elif "provider" == "azure":
+                    self.agent_settings["embedder"] = "azure"
+                elif "provider" == "openai":
+                    self.agent_settings["embedder"] = "openai"
+                elif "provider" == "palm":
+                    self.agent_settings["embedder"] = "google_vertex"
+        try:
             self.embedder = self.embedder_settings[self.agent_settings["embedder"]][
                 "embed"
             ]
             self.chunk_size = self.embedder_settings[self.agent_settings["embedder"]][
                 "chunk_size"
             ]
+        except:
+            self.embedder = self.default_embedder
+            self.chunk_size = 256
 
     def get_embedder_settings(self):
         if "API_URI" in self.agent_settings:
@@ -48,6 +58,25 @@ def get_embedder_settings(self):
                 "chunk_size": 256,
                 "embed": self.default_embedder,
             },
+            "local": {
+                "chunk_size": 1000,
+                "params": [
+                    "LOCAL_API_KEY",
+                    "AI_MODEL",
+                    "API_URI",
+                ],
+                "embed": embedding_functions.OpenAIEmbeddingFunction(
+                    api_key=self.agent_settings["LOCAL_API_KEY"]
+                    if "LOCAL_API_KEY" in self.agent_settings
+                    else "",
+                    model_name=self.agent_settings["AI_MODEL"]
+                    if "AI_MODEL" in self.agent_settings
+                    else "Mistral-7B-OpenOrca",
+                    api_base=self.agent_settings["API_URI"]
+                    if api_base
+                    else "https://localhost:8091/v1",
+                ),
+            },
             "azure": {
                 "chunk_size": 1000,
                 "params": [
@@ -88,15 +117,6 @@ def get_embedder_settings(self):
                 and "GOOGLE_API_KEY" in self.agent_settings
                 else self.default_embedder,
             },
-            "cohere": {
-                "chunk_size": 500,
-                "params": ["COHERE_API_KEY"],
-                "embed": embedding_functions.CohereEmbeddingFunction(
-                    api_key=self.agent_settings["COHERE_API_KEY"]
-                )
-                if "COHERE_API_KEY" in self.agent_settings
-                else self.default_embedder,
-            },
         }
         return embedder_settings
 

diff --git a/agixt/providers/local.py b/agixt/providers/local.py
@@ -32,7 +32,7 @@ def __init__(
         openai.api_key = LOCAL_API_KEY
 
     def models(self):
-        models = requests.get("http://localhost:8091/v1/models")
+        models = requests.get(f"{self.API_URI}/models")
         return models.json()
 
     async def instruct(self, prompt, tokens: int = 0):