Merge pull request #20 from Shreypandey/main

Embedding generation in LLMManager

jb-manager-bot/jb_manager_bot/parsers/utils.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 from openai import OpenAI, AzureOpenAI
 
 
@@ -119,3 +120,35 @@ def fn(cls, name, description, params, required):
                 },
             },
         }
+
+    @classmethod
+    def generate_embeddings(
+        cls,
+        inputs: List[str],
+        model: str,
+        openai_api_key=None,
+        azure_openai_api_key=None,
+        azure_openai_api_version=None,
+        azure_endpoint=None,
+        **kwargs
+    ):
+        """Use the OpenAI Embeddings API to generate embeddings for the given inputs."""
+        client = cls.get_client(
+            openai_api_key=openai_api_key,
+            azure_openai_api_key=azure_openai_api_key,
+            azure_openai_api_version=azure_openai_api_version,
+            azure_endpoint=azure_endpoint,
+        )
+        args = {
+            k: v
+            for k, v in kwargs.items()
+            if k
+            in [
+                "encoding_format",
+                "dimensions",
+            ]
+        }
+        response = client.embeddings.create(model=model, input=inputs, **args)
+        embeddings = [emb.embedding for emb in response.data]
+
+        return embeddings

jb-manager-bot/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jb-manager-bot"
-version = "0.1.4"
+version = "0.1.5"
 description = "Bot for JugalBandi Manager"
 authors = ["Shrey Pandey <[email protected]>", "Sameer Segal <[email protected]>", "Atharv Kirtikar <[email protected]>"]
 readme = "README.md"