lionagi 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

lionagi/integrations/provider/oai.py

@@ -1,8 +1,46 @@
+"""
+Copyright 2024 HaiyangLi
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 from os import getenv
 from lionagi.integrations.config.oai_configs import oai_schema
 from lionagi.libs.ln_api import BaseService, PayloadPackage
 
 
+allowed_kwargs = [
+    "model",
+    "frequency_penalty",
+    "n",
+    "presence_penalty",
+    "response_format",
+    "temperature",
+    "top_p",
+    "seed",
+    "stop",
+    "stream",
+    "stream_options",
+    "tools",
+    "tool_choice",
+    "user",
+    "max_tokens",
+    "logprobs",
+    "top_logprobs",
+    "logit_bias",
+]
+
+
 class OpenAIService(BaseService):
     """
     A service to interact with OpenAI's API endpoints.
@@ -30,6 +68,7 @@ class OpenAIService(BaseService):
         "audio_speech",
         "audio_transcriptions",
         "audio_translations",
+        "embeddings",
     ]
     schema = oai_schema
     key_scheme = "OPENAI_API_KEY"
@@ -51,6 +90,7 @@ class OpenAIService(BaseService):
             **kwargs,
         )
         self.active_endpoint = []
+        self.allowed_kwargs = allowed_kwargs
 
     async def serve(self, input_, endpoint="chat/completions", method="post", **kwargs):
         """
@@ -84,7 +124,7 @@ class OpenAIService(BaseService):
         else:
             return ValueError(f"{endpoint} is currently not supported")
 
-    async def serve_chat(self, messages, **kwargs):
+    async def serve_chat(self, messages, required_tokens=None, **kwargs):
         """
         Serves the chat completion request with the given messages.
 
@@ -101,16 +141,63 @@ class OpenAIService(BaseService):
         if "chat/completions" not in self.active_endpoint:
             await self.init_endpoint("chat/completions")
             self.active_endpoint.append("chat/completions")
+
+        msgs = []
+
+        for msg in messages:
+            if isinstance(msg, dict):
+                content = msg.get("content")
+                if isinstance(content, (dict, str)):
+                    msgs.append({"role": msg["role"], "content": content})
+                elif isinstance(content, list):
+                    _content = []
+                    for i in content:
+                        if "text" in i:
+                            _content.append({"type": "text", "text": str(i["text"])})
+                        elif "image_url" in i:
+                            _content.append(
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": f"{i['image_url'].get('url')}",
+                                        "detail": i["image_url"].get("detail", "low"),
+                                    },
+                                }
+                            )
+                    msgs.append({"role": msg["role"], "content": _content})
+
         payload = PayloadPackage.chat_completion(
-            messages,
+            msgs,
             self.endpoints["chat/completions"].config,
             self.schema["chat/completions"],
             **kwargs,
         )
-
         try:
-            completion = await self.call_api(payload, "chat/completions", "post")
+            completion = await self.call_api(
+                payload, "chat/completions", "post", required_tokens=required_tokens
+            )
             return payload, completion
         except Exception as e:
             self.status_tracker.num_tasks_failed += 1
             raise e
+
+    async def serve_embedding(self, embed_str, required_tokens=None, **kwargs):
+        if "embeddings" not in self.active_endpoint:
+            await self.init_endpoint("embeddings")
+            self.active_endpoint.append("embeddings")
+
+        payload = PayloadPackage.embeddings(
+            embed_str,
+            self.endpoints["embeddings"].config,
+            self.schema["embeddings"],
+            **kwargs,
+        )
+
+        try:
+            embed = await self.call_api(
+                payload, "embeddings", "post", required_tokens=required_tokens
+            )
+            return payload, embed
+        except Exception as e:
+            self.status_tracker.num_tasks_failed += 1
+            raise e

lionagi/integrations/provider/ollama.py

@@ -4,7 +4,7 @@ from lionagi.integrations.config.ollama_configs import model
 allowed_kwargs = [
     "model",
     "frequency_penalty",
-    "max_tokens",
+    # "max_tokens",
     "n",
     "presence_penalty",
     "response_format",
@@ -32,7 +32,8 @@ class OllamaService(BaseService):
 
         self.ollama = ollama
         self.model = model
-        self.client = self.ollama.AsyncClient(**kwargs)
+        self.client = self.ollama.AsyncClient()
+        self.allowed_kwargs = allowed_kwargs
 
     async def serve_chat(self, messages, **kwargs):
         config = {}
@@ -42,11 +43,11 @@ class OllamaService(BaseService):
 
         self.ollama.pull(self.model)
         payload = {"messages": messages}
+        if "model" not in config:
+            config["model"] = self.model
 
         try:
-            completion = await self.client.chat(
-                model=self.model, messages=messages, **config
-            )
+            completion = await self.client.chat(messages=messages, **config)
             completion["choices"] = [{"message": completion.pop("message")}]
             return payload, completion
         except Exception as e:

lionagi/integrations/provider/openrouter.py

@@ -2,6 +2,27 @@ from os import getenv
 from lionagi.integrations.config.openrouter_configs import openrouter_schema
 from lionagi.libs.ln_api import BaseService, PayloadPackage
 
+allowed_kwargs = [
+    "model",
+    "frequency_penalty",
+    "n",
+    "presence_penalty",
+    "response_format",
+    "temperature",
+    "top_p",
+    "seed",
+    "stop",
+    "stream",
+    "stream_options",
+    "tools",
+    "tool_choice",
+    "user",
+    "max_tokens",
+    "logprobs",
+    "top_logprobs",
+    "logit_bias",
+]
+
 
 class OpenRouterService(BaseService):
     base_url = "https://openrouter.ai/api/v1/"
@@ -26,8 +47,33 @@ class OpenRouterService(BaseService):
             **kwargs,
         )
         self.active_endpoint = []
+        self.allowed_kwargs = allowed_kwargs
 
     async def serve(self, input_, endpoint="chat/completions", method="post", **kwargs):
+        """
+        Serves the input using the specified endpoint and method.
+
+        Args:
+                input_: The input text to be processed.
+                endpoint: The API endpoint to use for processing.
+                method: The HTTP method to use for the request.
+                **kwargs: Additional keyword arguments to pass to the payload creation.
+
+        Returns:
+                A tuple containing the payload and the completion assistant_response from the API.
+
+        Raises:
+                ValueError: If the specified endpoint is not supported.
+
+        Examples:
+                >>> service = OpenAIService(api_key="your_api_key")
+                >>> asyncio.run(service.serve("Hello, world!","chat/completions"))
+                (payload, completion)
+
+                >>> service = OpenAIService()
+                >>> asyncio.run(service.serve("Convert this text to speech.","audio_speech"))
+                ValueError: 'audio_speech' is currently not supported
+        """
         if endpoint not in self.active_endpoint:
             await self.init_endpoint(endpoint)
         if endpoint == "chat/completions":
@@ -35,19 +81,80 @@ class OpenRouterService(BaseService):
         else:
             return ValueError(f"{endpoint} is currently not supported")
 
-    async def serve_chat(self, messages, **kwargs):
-        endpoint = "chat/completions"
+    async def serve_chat(self, messages, required_tokens=None, **kwargs):
+        """
+        Serves the chat completion request with the given messages.
+
+        Args:
+                messages: The messages to be included in the chat completion.
+                **kwargs: Additional keyword arguments for payload creation.
+
+        Returns:
+                A tuple containing the payload and the completion assistant_response from the API.
+
+        Raises:
+                Exception: If the API call fails.
+        """
+        if "chat/completions" not in self.active_endpoint:
+            await self.init_endpoint("chat/completions")
+            self.active_endpoint.append("chat/completions")
+
+        msgs = []
+
+        for msg in messages:
+            if isinstance(msg, dict):
+                content = msg.get("content")
+                if isinstance(content, (dict, str)):
+                    msgs.append({"role": msg["role"], "content": content})
+                elif isinstance(content, list):
+                    _content = []
+                    for i in content:
+                        if "text" in i:
+                            _content.append({"type": "text", "text": str(i["text"])})
+                        elif "image_url" in i:
+                            _content.append(
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": f"{i['image_url'].get('url')}",
+                                        "detail": i["image_url"].get("detail", "low"),
+                                    },
+                                }
+                            )
+                    msgs.append({"role": msg["role"], "content": _content})
 
-        if endpoint not in self.active_endpoint:
-            await self.init_endpoint(endpoint)
-            self.active_endpoint.append(endpoint)
         payload = PayloadPackage.chat_completion(
-            messages, self.endpoints[endpoint].config, self.schema[endpoint], **kwargs
+            msgs,
+            self.endpoints["chat/completions"].config,
+            self.schema["chat/completions"],
+            **kwargs,
         )
-
         try:
-            completion = await self.call_api(payload, endpoint, "post")
+            completion = await self.call_api(
+                payload, "chat/completions", "post", required_tokens=required_tokens
+            )
             return payload, completion
         except Exception as e:
             self.status_tracker.num_tasks_failed += 1
             raise e
+
+    # async def serve_embedding(self, embed_str, required_tokens=None, **kwargs):
+    #     if "embeddings" not in self.active_endpoint:
+    #         await self.init_endpoint("embeddings")
+    #         self.active_endpoint.append("embeddings")
+
+    #     payload = PayloadPackage.embeddings(
+    #         embed_str,
+    #         self.endpoints["embeddings"].config,
+    #         self.schema["embeddings"],
+    #         **kwargs,
+    #     )
+
+    #     try:
+    #         embed = await self.call_api(
+    #             payload, "embeddings", "post", required_tokens=required_tokens
+    #         )
+    #         return payload, embed
+    #     except Exception as e:
+    #         self.status_tracker.num_tasks_failed += 1
+    #         raise e

lionagi/integrations/provider/services.py

@@ -131,6 +131,6 @@ class Services:
                 kwargs (Optional[Any]): additional kwargs for calling the model
         """
 
-        from lionagi.integrations.provider.mlx_service import MlXService
+        from lionagi.integrations.provider.mlx_service import MLXService
 
-        return MlXService(**kwargs)
+        return MLXService(**kwargs)

lionagi/integrations/provider/transformers.py

@@ -5,7 +5,7 @@ from lionagi.libs.sys_util import SysUtil
 from lionagi.libs.ln_api import BaseService
 
 allowed_kwargs = [
-    "model",
+    # "model",
     "tokenizer",
     "modelcard",
     "framework",
@@ -17,6 +17,9 @@ allowed_kwargs = [
     "torch_dtype",
     "min_length_for_response",
     "minimum_tokens",
+    "mask_token",
+    "max_length",
+    "max_new_tokens",
 ]
 
 
@@ -52,6 +55,7 @@ class TransformersService(BaseService):
         self.task = task
         self.model = model
         self.config = config
+        self.allowed_kwargs = allowed_kwargs
         try:
             from transformers import pipeline
 
@@ -59,19 +63,11 @@ class TransformersService(BaseService):
         except ImportError:
             try:
                 if not SysUtil.is_package_installed("torch"):
-                    in_ = input(
-                        "PyTorch is required for transformers. Would you like to install it now? (y/n): "
-                    )
-                    if in_ == "y":
-                        install_pytorch()
+                    install_pytorch()
                 if not SysUtil.is_package_installed("transformers"):
-                    in_ = input(
-                        "transformers is required. Would you like to install it now? (y/n): "
+                    SysUtil.install_import(
+                        package_name="transformers", import_name="pipeline"
                     )
-                    if in_ == "y":
-                        SysUtil.install_import(
-                            package_name="transformers", import_name="pipeline"
-                        )
                     from transformers import pipeline
 
                     self.pipeline = pipeline
@@ -92,19 +88,19 @@ class TransformersService(BaseService):
         payload = {"messages": messages}
         config = {}
         for k, v in kwargs.items():
+            if k == "max_tokens":
+                config["max_new_tokens"] = v
             if k in allowed_kwargs:
                 config[k] = v
 
-        conversation = self.pipe(str(messages), **config)
-
-        texts = conversation[-1]["generated_text"]
-        msgs = (
-            str(texts.split("]")[1:])
-            .replace("\\n", "")
-            .replace("['", "")
-            .replace("\\", "")
-        )
+        msg = "".join([i["content"] for i in messages if i["role"] == "user"])
+        conversation = ""
+        response = self.pipe(msg, **config)
+        try:
+            conversation = response[0]["generated_text"]
+        except:
+            conversation = response
 
-        completion = {"model": self.pipe.model, "choices": [{"message": msgs}]}
+        completion = {"choices": [{"message": {"content": conversation}}]}
 
         return payload, completion

lionagi/integrations/storage/__init__.py

@@ -1,3 +1,3 @@
-from lionagi.integrations.storage.neo4j import Neo4j
-
-__all__ = ["Neo4j"]
+# from lionagi.integrations.storage.neo4j import Neo4j
+#
+# __all__ = ["Neo4j"]