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

lionagi/integrations/loader/load_util.py

@@ -1,10 +1,10 @@
 # use utils and schema
-import math
 from enum import Enum
 from pathlib import Path
 from typing import List, Union, Dict, Any, Tuple
 
 from lionagi.libs import convert, func_call
+from lionagi.libs.ln_tokenize import TokenizeUtil
 from lionagi.core.generic import Node
 
 
@@ -57,7 +57,7 @@ def dir_to_path(
 
 
 def dir_to_nodes(
-    dir: str,
+    dir_: str,
     ext: Union[List[str], str],
     recursive: bool = False,
     flatten: bool = True,
@@ -84,82 +84,11 @@ def dir_to_nodes(
         # converting them into Node objects.
     """
 
-    path_list = dir_to_path(dir, ext, recursive, flatten)
+    path_list = dir_to_path(dir=dir_, ext=ext, recursive=recursive, flatten=flatten)
     files_info = func_call.lcall(path_list, read_text, clean=clean_text)
     return func_call.lcall(files_info, lambda x: Node(content=x[0], metadata=x[1]))
 
 
-def chunk_text(
-    input: str, chunk_size: int, overlap: float, threshold: int
-) -> List[Union[str, None]]:
-    """
-    Chunks the input text into smaller parts, with optional overlap and threshold for final chunk.
-
-    Parameters:
-        input (str): The input text to chunk.
-
-        chunk_size (int): The size of each chunk.
-
-        overlap (float): The amount of overlap between chunks.
-
-        threshold (int): The minimum size of the final chunk.
-
-    Returns:
-        List[Union[str, None]]: A list of text chunks.
-
-    Raises:
-        ValueError: If an error occurs during chunking.
-    """
-
-    def _chunk_n1():
-        return [input]
-
-    def _chunk_n2():
-        chunks = []
-        chunks.append(input[: chunk_size + overlap_size])
-
-        if len(input) - chunk_size > threshold:
-            chunks.append(input[chunk_size - overlap_size :])
-        else:
-            return _chunk_n1()
-
-        return chunks
-
-    def _chunk_n3():
-        chunks = []
-        chunks.append(input[: chunk_size + overlap_size])
-        for i in range(1, n_chunks - 1):
-            start_idx = chunk_size * i - overlap_size
-            end_idx = chunk_size * (i + 1) + overlap_size
-            chunks.append(input[start_idx:end_idx])
-
-        if len(input) - chunk_size * (n_chunks - 1) > threshold:
-            chunks.append(input[chunk_size * (n_chunks - 1) - overlap_size :])
-        else:
-            chunks[-1] += input[chunk_size * (n_chunks - 1) + overlap_size :]
-
-        return chunks
-
-    try:
-        if not isinstance(input, str):
-            input = convert.to_str(input)
-
-        n_chunks = math.ceil(len(input) / chunk_size)
-        overlap_size = int(overlap / 2)
-
-        if n_chunks == 1:
-            return _chunk_n1()
-
-        elif n_chunks == 2:
-            return _chunk_n2()
-
-        elif n_chunks > 2:
-            return _chunk_n3()
-
-    except Exception as e:
-        raise ValueError(f"An error occurred while chunking the text. {e}")
-
-
 def read_text(filepath: str, clean: bool = True) -> Tuple[str, dict]:
     """
     Reads text from a file and optionally cleans it, returning the content and metadata.
@@ -202,9 +131,9 @@ def read_text(filepath: str, clean: bool = True) -> Tuple[str, dict]:
             content = f.read()
             if clean:
                 # Define characters to replace and their replacements
-                replacements = {"\\": " ", "\n": " ", "\t": " ", "  ": " ", "'": " "}
+                replacements = {"\\": "", "\n\n": "\n"}
                 for old, new in replacements.items():
-                    content = content.replace(old, new)
+                    content = content.replace(old, new).strip()
             metadata = _get_metadata()
             return content, metadata
     except Exception as e:
@@ -223,7 +152,7 @@ def _file_to_chunks(
             "chunk_overlap": overlap,
             "chunk_threshold": threshold,
         }
-        chunks = chunk_text(
+        chunks = TokenizeUtil.chunk_by_chars(
             input[field], chunk_size=chunk_size, overlap=overlap, threshold=threshold
         )
         logs = []

lionagi/integrations/provider/_mapping.py

@@ -0,0 +1,46 @@
+from .oai import OpenAIService
+from .openrouter import OpenRouterService
+from .ollama import OllamaService
+from .transformers import TransformersService
+from .litellm import LiteLLMService
+from .mlx_service import MLXService
+from lionagi.integrations.config.oai_configs import oai_schema
+from lionagi.integrations.config.openrouter_configs import openrouter_schema
+
+SERVICE_PROVIDERS_MAPPING = {
+    "openai": {
+        "service": OpenAIService,
+        "schema": oai_schema,
+        "default_model": "gpt-3.5-turbo",
+    },
+    "openrouter": {
+        "service": OpenRouterService,
+        "schema": openrouter_schema,
+        "default_model": "gpt-3.5-turbo",
+    },
+    "litellm": {
+        "service": LiteLLMService,
+        "schema": oai_schema,
+        "default_model": "gpt-3.5-turbo",
+    },
+    "ollama": {
+        "service": OllamaService,
+        "schema": {"model": "llama3"},
+        "default_model": "llama3",
+    },
+    "transformers": {
+        "service": TransformersService,
+        "schema": {"model": "gpt2"},
+        "default_model": "gpt2",
+    },
+    "mlx": {
+        "service": MLXService,
+        "schema": {"model": "mlx-community/OLMo-7B-hf-4bit-mlx"},
+        "default_model": "mlx-community/OLMo-7B-hf-4bit-mlx",
+    },
+}
+
+# TODO
+# "Ollama": OllamaService,
+# "Transformers": TransformersService,
+# "MLX": MLXService,

lionagi/integrations/provider/litellm.py

@@ -31,12 +31,13 @@ class LiteLLMService(BaseService):
         self.acompletion = acompletion
         self.model = model
         self.kwargs = kwargs
+        self.allowed_kwargs = allowed_kwargs
 
     async def serve_chat(self, messages, **kwargs):
         payload = {"messages": messages}
         config = {}
         for k, v in kwargs.items():
-            if k in allowed_kwargs:
+            if k in self.allowed_kwargs:
                 config[k] = v
 
         kwargs = {**self.kwargs, **config}

lionagi/integrations/provider/mlx_service.py

@@ -1,13 +1,18 @@
+import re
 from lionagi.libs.sys_util import SysUtil
 import lionagi.libs.ln_convert as convert
 from lionagi.libs.ln_api import BaseService
 from lionagi.integrations.config.mlx_configs import model
 
 
-class MlXService(BaseService):
+class MLXService(BaseService):
     def __init__(self, model=model, **kwargs):
 
         SysUtil.check_import("mlx_lm")
+        SysUtil.check_import("ipywidgets")
+
+        if model is not None and "olmo" in str(model).lower():
+            SysUtil.check_import("olmo", pip_name="ai2-olmo")
 
         from mlx_lm import load, generate
 
@@ -19,16 +24,13 @@ class MlXService(BaseService):
         self.model = model_
         self.tokenizer = tokenizer
         self.generate = generate
+        self.allowed_kwargs = []
 
     async def serve_chat(self, messages, **kwargs):
         if "verbose" not in kwargs.keys():
-            verbose = True
+            verbose = False
 
-        prompts = [
-            convert.to_dict(msg["content"])["instruction"]
-            for msg in messages
-            if msg["role"] == "user"
-        ]
+        prompts = [msg["content"] for msg in messages if msg["role"] == "user"]
 
         payload = {"messages": messages}
 
@@ -39,8 +41,13 @@ class MlXService(BaseService):
                 prompt=f"{prompts[-1]} \nOutput: ",
                 verbose=verbose,
             )
-            completion = {"model": self.model_name, "choices": [{"message": response}]}
-
+            if "```" in response:
+                regex = re.compile(r"```[\s\S]*?```")
+                matches = regex.findall(response)
+                msg = matches[0].strip("```")
+                completion = {"choices": [{"message": {"content": msg}}]}
+            else:
+                completion = {"choices": [{"message": {"content": response}}]}
             return payload, completion
         except Exception as e:
             self.status_tracker.num_tasks_failed += 1

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"]