lionagi 0.0.316__py3-none-any.whl → 0.1.1__py3-none-any.whl

lionagi/libs/ln_knowledge_graph.py

@@ -0,0 +1,405 @@
+import math
+from lionagi.libs import CallDecorator as cd
+
+
+class KnowledgeBase:
+    """
+    A class to represent a Knowledge Base (KB) containing entities, relations, and sources.
+
+    Attributes:
+        entities (dict): A dictionary of entities in the KB, where the keys are entity titles, and the values are
+                         entity information (excluding the title).
+        relations (list): A list of relations in the KB, where each relation is a dictionary containing information
+                          about the relation (head, type, tail) and metadata (article_url and spans).
+        sources (dict): A dictionary of information about the sources of relations, where the keys are article URLs,
+                       and the values are source data (article_title and article_publish_date).
+
+    Methods:
+        merge_with_kb(kb2): Merge another Knowledge Base (kb2) into this KB.
+        are_relations_equal(r1, r2): Check if two relations (r1 and r2) are equal.
+        exists_relation(r1): Check if a relation (r1) already exists in the KB.
+        merge_relations(r2): Merge the information from relation r2 into an existing relation in the KB.
+        get_wikipedia_data(candidate_entity): Get data for a candidate entity from Wikipedia.
+        add_entity(e): Add an entity to the KB.
+        add_relation(r, article_title, article_publish_date): Add a relation to the KB.
+        print(): Print the entities, relations, and sources in the KB.
+        extract_relations_from_model_output(text): Extract relations from the model output text.
+
+    """
+
+    def __init__(self):
+        """
+        Initialize an empty Knowledge Base (KB) with empty dictionaries for entities, relations, and sources.
+        """
+        self.entities = {}  # { entity_title: {...} }
+        self.relations = []  # [ head: entity_title, type: ..., tail: entity_title,
+        # meta: { article_url: { spans: [...] } } ]
+        self.sources = {}  # { article_url: {...} }
+
+    def merge_with_kb(self, kb2):
+        """
+        Merge another Knowledge Base (KB) into this KB.
+
+        Args:
+            kb2 (KnowledgeBase): The Knowledge Base (KB) to merge into this KB.
+        """
+        for r in kb2.relations:
+            article_url = list(r["meta"].keys())[0]
+            source_data = kb2.sources[article_url]
+            self.add_relation(
+                r, source_data["article_title"], source_data["article_publish_date"]
+            )
+
+    def are_relations_equal(self, r1, r2):
+        """
+        Check if two relations (r1 and r2) are equal.
+
+        Args:
+            r1 (dict): The first relation to compare.
+            r2 (dict): The second relation to compare.
+
+        Returns:
+            bool: True if the relations are equal, False otherwise.
+        """
+        return all(r1[attr] == r2[attr] for attr in ["head", "type", "tail"])
+
+    def exists_relation(self, r1):
+        """
+        Check if a relation (r1) already exists in the KB.
+
+        Args:
+            r1 (dict): The relation to check for existence in the KB.
+
+        Returns:
+            bool: True if the relation exists in the KB, False otherwise.
+        """
+        return any(self.are_relations_equal(r1, r2) for r2 in self.relations)
+
+    def merge_relations(self, r2):
+        """
+        Merge the information from relation r2 into an existing relation in the KB.
+
+        Args:
+            r2 (dict): The relation to merge into an existing relation in the KB.
+        """
+        r1 = [r for r in self.relations if self.are_relations_equal(r2, r)][0]
+
+        # if different article
+        article_url = list(r2["meta"].keys())[0]
+        if article_url not in r1["meta"]:
+            r1["meta"][article_url] = r2["meta"][article_url]
+
+        # if existing article
+        else:
+            spans_to_add = [
+                span
+                for span in r2["meta"][article_url]["spans"]
+                if span not in r1["meta"][article_url]["spans"]
+            ]
+            r1["meta"][article_url]["spans"] += spans_to_add
+
+    @cd.cache(maxsize=10000)
+    def get_wikipedia_data(self, candidate_entity):
+        """
+        Get data for a candidate entity from Wikipedia.
+
+        Args:
+            candidate_entity (str): The candidate entity title.
+
+        Returns:
+            dict: A dictionary containing information about the candidate entity (title, url, summary).
+                  None if the entity does not exist in Wikipedia.
+        """
+        try:
+            from lionagi.libs import SysUtil
+
+            SysUtil.check_import("wikipedia")
+            import wikipedia  # type: ignore
+        except Exception as e:
+            raise Exception("wikipedia package is not installed {e}")
+
+        try:
+            page = wikipedia.page(candidate_entity, auto_suggest=False)
+            entity_data = {
+                "title": page.title,
+                "url": page.url,
+                "summary": page.summary,
+            }
+            return entity_data
+        except:
+            return None
+
+    def add_entity(self, e):
+        """
+        Add an entity to the KB.
+
+        Args:
+            e (dict): A dictionary containing information about the entity (title and additional attributes).
+        """
+        self.entities[e["title"]] = {k: v for k, v in e.items() if k != "title"}
+
+    def add_relation(self, r, article_title, article_publish_date):
+        """
+        Add a relation to the KB.
+
+        Args:
+            r (dict): A dictionary containing information about the relation (head, type, tail, and metadata).
+            article_title (str): The title of the article containing the relation.
+            article_publish_date (str): The publish date of the article.
+        """
+        # check on wikipedia
+        candidate_entities = [r["head"], r["tail"]]
+        entities = [self.get_wikipedia_data(ent) for ent in candidate_entities]
+
+        # if one entity does not exist, stop
+        if any(ent is None for ent in entities):
+            return
+
+        # manage new entities
+        for e in entities:
+            self.add_entity(e)
+
+        # rename relation entities with their wikipedia titles
+        r["head"] = entities[0]["title"]
+        r["tail"] = entities[1]["title"]
+
+        # add source if not in kb
+        article_url = list(r["meta"].keys())[0]
+        if article_url not in self.sources:
+            self.sources[article_url] = {
+                "article_title": article_title,
+                "article_publish_date": article_publish_date,
+            }
+
+        # manage new relation
+        if not self.exists_relation(r):
+            self.relations.append(r)
+        else:
+            self.merge_relations(r)
+
+    def print(self):
+        """
+        Print the entities, relations, and sources in the KB.
+
+        Returns:
+            None
+        """
+        print("Entities:")
+        for e in self.entities.items():
+            print(f"  {e}")
+        print("Relations:")
+        for r in self.relations:
+            print(f"  {r}")
+        print("Sources:")
+        for s in self.sources.items():
+            print(f"  {s}")
+
+    @staticmethod
+    def extract_relations_from_model_output(text):
+        """
+        Extract relations from the model output text.
+
+        Args:
+            text (str): The model output text containing relations.
+
+        Returns:
+            list: A list of dictionaries, where each dictionary represents a relation (head, type, tail).
+        """
+        relations = []
+        relation, subject, relation, object_ = "", "", "", ""
+        text = text.strip()
+        current = "x"
+        text_replaced = text.replace("<s>", "").replace("<pad>", "").replace("</s>", "")
+        for token in text_replaced.split():
+            if token == "<triplet>":
+                current = "t"
+                if relation != "":
+                    relations.append(
+                        {
+                            "head": subject.strip(),
+                            "type": relation.strip(),
+                            "tail": object_.strip(),
+                        }
+                    )
+                    relation = ""
+                subject = ""
+            elif token == "<subj>":
+                current = "s"
+                if relation != "":
+                    relations.append(
+                        {
+                            "head": subject.strip(),
+                            "type": relation.strip(),
+                            "tail": object_.strip(),
+                        }
+                    )
+                object_ = ""
+            elif token == "<obj>":
+                current = "o"
+                relation = ""
+            else:
+                if current == "t":
+                    subject += " " + token
+                elif current == "s":
+                    object_ += " " + token
+                elif current == "o":
+                    relation += " " + token
+        if subject != "" and relation != "" and object_ != "":
+            relations.append(
+                {
+                    "head": subject.strip(),
+                    "type": relation.strip(),
+                    "tail": object_.strip(),
+                }
+            )
+        return relations
+
+
+class KGTripletExtractor:
+    """
+    A class to perform knowledge graph triplet extraction from text using a pre-trained model.
+
+    Methods:
+        text_to_wiki_kb(text, model=None, tokenizer=None, device='cpu', span_length=512,
+                        article_title=None, article_publish_date=None, verbose=False):
+            Extract knowledge graph triplets from text and create a KnowledgeBase (KB) containing entities and relations.
+
+    """
+
+    @staticmethod
+    def text_to_wiki_kb(
+        text,
+        model=None,
+        tokenizer=None,
+        device="cpu",
+        span_length=512,
+        article_title=None,
+        article_publish_date=None,
+        verbose=False,
+    ):
+        from lionagi.integrations.bridge.transformers_.install_ import (
+            install_transformers,
+        )
+
+        try:
+            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer  # type: ignore
+        except ImportError:
+            install_transformers()
+        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer  # type: ignore
+        import torch  # type: ignore
+
+        """
+        Extract knowledge graph triplets from text and create a KnowledgeBase (KB) containing entities and relations.
+
+        Args:
+            text (str): The input text from which triplets will be extracted.
+            model (AutoModelForSeq2SeqLM, optional): The pre-trained model for triplet extraction. Defaults to None.
+            tokenizer (AutoTokenizer, optional): The tokenizer for the model. Defaults to None.
+            device (str, optional): The device to run the model on (e.g., 'cpu', 'cuda'). Defaults to 'cpu'.
+            span_length (int, optional): The maximum span length for input text segmentation. Defaults to 512.
+            article_title (str, optional): The title of the article containing the input text. Defaults to None.
+            article_publish_date (str, optional): The publish date of the article. Defaults to None.
+            verbose (bool, optional): Whether to enable verbose mode for debugging. Defaults to False.
+
+        Returns:
+            KnowledgeBase: A KnowledgeBase (KB) containing extracted entities, relations, and sources.
+
+        """
+
+        if not any([model, tokenizer]):
+            tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
+            model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
+            model.to(device)
+
+        inputs = tokenizer([text], return_tensors="pt")
+
+        num_tokens = len(inputs["input_ids"][0])
+        if verbose:
+            print(f"Input has {num_tokens} tokens")
+        num_spans = math.ceil(num_tokens / span_length)
+        if verbose:
+            print(f"Input has {num_spans} spans")
+        overlap = math.ceil(
+            (num_spans * span_length - num_tokens) / max(num_spans - 1, 1)
+        )
+        spans_boundaries = []
+        start = 0
+        for i in range(num_spans):
+            spans_boundaries.append(
+                [start + span_length * i, start + span_length * (i + 1)]
+            )
+            start -= overlap
+        if verbose:
+            print(f"Span boundaries are {spans_boundaries}")
+
+        # transform input with spans
+        tensor_ids = [
+            inputs["input_ids"][0][boundary[0] : boundary[1]]
+            for boundary in spans_boundaries
+        ]
+        tensor_masks = [
+            inputs["attention_mask"][0][boundary[0] : boundary[1]]
+            for boundary in spans_boundaries
+        ]
+
+        inputs = {
+            "input_ids": torch.stack(tensor_ids).to(device),
+            "attention_mask": torch.stack(tensor_masks).to(device),
+        }
+
+        # generate relations
+        num_return_sequences = 3
+        gen_kwargs = {
+            "max_length": 512,
+            "length_penalty": 0,
+            "num_beams": 3,
+            "num_return_sequences": num_return_sequences,
+        }
+        generated_tokens = model.generate(
+            **inputs,
+            **gen_kwargs,
+        )
+
+        # decode relations
+        decoded_preds = tokenizer.batch_decode(
+            generated_tokens, skip_special_tokens=False
+        )
+
+        # create kb
+        kb = KnowledgeBase()
+        i = 0
+        for sentence_pred in decoded_preds:
+            current_span_index = i // num_return_sequences
+            relations = KnowledgeBase.extract_relations_from_model_output(sentence_pred)
+            for relation in relations:
+                relation["meta"] = {
+                    "article_url": {"spans": [spans_boundaries[current_span_index]]}
+                }
+                kb.add_relation(relation, article_title, article_publish_date)
+            i += 1
+        return kb
+
+
+class KGraph:
+    """
+    A class representing a Knowledge Graph (KGraph) for extracting relations from text.
+
+    Methods:
+        text_to_wiki_kb(text, model=None, tokenizer=None, device='cpu', span_length=512, article_title=None,
+                        article_publish_date=None, verbose=False):
+            Extract relations from input text and create a Knowledge Base (KB) containing entities and relations.
+    """
+
+    @staticmethod
+    def text_to_wiki_kb(text, **kwargs):
+        """
+        Extract relations from input text and create a Knowledge Base (KB) containing entities and relations.
+
+        Args:
+            text (str): The input text from which relations are extracted.
+            **kwargs: Additional keyword arguments passed to the underlying extraction method.
+
+        Returns:
+            KnowledgeBase: A Knowledge Base (KB) containing entities and relations extracted from the input text.
+        """
+        return KGTripletExtractor.text_to_wiki_kb(text, **kwargs)

lionagi/libs/ln_parse.py

@@ -1,11 +1,14 @@
+from collections.abc import Callable
 import re
 import inspect
 import itertools
-from collections.abc import Callable
+import contextlib
+from functools import singledispatchmethod
 from typing import Any
 import numpy as np
 import lionagi.libs.ln_convert as convert
 
+
 md_json_char_map = {"\n": "\\n", "\r": "\\r", "\t": "\\t", '"': '\\"'}
 
 
@@ -590,14 +593,15 @@ class StringMatch:
         return d[m][n]
 
     @staticmethod
-    def correct_keys(output_fields, out_, score_func=None):
+    def correct_dict_keys(keys: dict | list[str], dict_, score_func=None):
         if score_func is None:
             score_func = StringMatch.jaro_winkler_similarity
-        fields_set = set(output_fields.keys())
+
+        fields_set = set(keys if isinstance(keys, list) else keys.keys())
         corrected_out = {}
         used_keys = set()
 
-        for k, v in out_.items():
+        for k, v in dict_.items():
             if k in fields_set:
                 corrected_out[k] = v
                 fields_set.remove(k)  # Remove the matched key
@@ -614,8 +618,8 @@ class StringMatch:
                 fields_set.remove(best_match)  # Remove the matched key
                 used_keys.add(best_match)
 
-        if len(used_keys) < len(out_):
-            for k, v in out_.items():
+        if len(used_keys) < len(dict_):
+            for k, v in dict_.items():
                 if k not in used_keys:
                     corrected_out[k] = v
 
@@ -637,3 +641,36 @@ class StringMatch:
         # Find the index of the highest score
         max_score_index = np.argmax(scores)
         return correct_words_list[max_score_index]
+
+    @staticmethod
+    def force_validate_dict(x, keys: dict | list[str]) -> dict:
+        out_ = x
+
+        if isinstance(out_, str):
+            # first try to parse it straight as a fuzzy json
+            try:
+                out_ = ParseUtil.fuzzy_parse_json(out_)
+            except Exception:
+                try:
+                    # if failed we try to extract the json block and parse it
+                    out_ = ParseUtil.md_to_json(out_)
+                except Exception:
+                    # if still failed we try to extract the json block using re and parse it again
+                    match = re.search(r"```json\n({.*?})\n```", out_, re.DOTALL)
+                    if match:
+                        out_ = match.group(1)
+                        try:
+                            out_ = ParseUtil.fuzzy_parse_json(out_)
+                        except:
+                            try:
+                                out_ = ParseUtil.fuzzy_parse_json(
+                                    out_.replace("'", '"')
+                                )
+                            except:
+                                pass
+
+        if isinstance(out_, dict):
+            try:
+                return StringMatch.correct_dict_keys(keys, out_)
+            except Exception as e:
+                raise ValueError(f"Failed to force_validate_dict for input: {x}") from e

lionagi/libs/ln_queue.py

@@ -0,0 +1,101 @@
+"""
+A class that manages asynchronous task processing with controlled concurrency.
+"""
+
+from typing import Any, Callable
+import asyncio
+from lionagi.libs import func_call
+
+
+class AsyncQueue:
+    """
+    This class handles the enqueueing and processing of tasks with a limit on
+    how many can run simultaneously, using an asyncio.Queue for task storage and
+    an asyncio.Semaphore to limit concurrency.
+
+    Attributes:
+        queue (asyncio.Queue): The queue to store tasks.
+        _stop_event (asyncio.Event): Event to signal processing should stop.
+        max_concurrent_tasks (int): Maximum number of tasks processed concurrently.
+        semaphore (asyncio.Semaphore): Controls concurrent access to task execution.
+    """
+
+    def __init__(self, max_concurrent_tasks=5):
+        """
+        Initializes the AsyncQueue with a concurrency limit.
+
+        Args:
+            max_concurrent_tasks (int): The maximum number of concurrent tasks
+                                        allowed. Default is 5.
+        """
+        self.queue = asyncio.Queue()
+        self._stop_event = asyncio.Event()
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.semaphore = asyncio.Semaphore(max_concurrent_tasks)
+
+    async def enqueue(self, input_: Any) -> None:
+        """
+        Enqueues an item to be processed asynchronously.
+
+        Args:
+            input_ (Any): The item to be enqueued.
+        """
+        await self.queue.put(input_)
+
+    async def dequeue(self) -> Any:
+        """
+        Dequeues an item for processing.
+
+        Returns:
+            Any: The dequeued item.
+        """
+        return await self.queue.get()
+
+    async def join(self) -> None:
+        """Waits for all items in the queue to be processed."""
+        await self.queue.join()
+
+    async def stop(self) -> None:
+        """Signals the queue to stop processing new items."""
+        self._stop_event.set()
+
+    def stopped(self) -> bool:
+        """
+        Checks if the stop signal has been issued.
+
+        Returns:
+            bool: True if the queue has been stopped, otherwise False.
+        """
+        return self._stop_event.is_set()
+
+    async def process_requests(self, func: Callable, retry_kwargs: dict = {}) -> None:
+        """
+        Processes tasks from the queue using the provided function with retries.
+
+        This method continuously processes tasks from the queue using the specified
+        function until a stop event is triggered. Handles concurrency using a
+        semaphore and manages task completion.
+
+        Args:
+            func (Callable): The function to process each task.
+            retry_kwargs (dict): Keyword arguments for retry behavior. Default is
+                                 an empty dictionary.
+        """
+        tasks = set()
+        while not self.stopped():
+            if len(tasks) >= self.max_concurrent_tasks:
+                _, done = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+                tasks.difference_update(done)
+
+            async with self.semaphore:
+                input_ = await self.dequeue()
+                if input_ is None:
+                    await self.stop()
+                    break
+                task = asyncio.create_task(
+                    func_call.rcall(func, input_, **retry_kwargs)
+                )
+                tasks.add(task)
+
+        if tasks:
+            await asyncio.wait(tasks)

lionagi/libs/ln_tokenizer.py

@@ -0,0 +1,57 @@
+import re
+
+
+class BaseToken:
+    def __init__(self, type_, value):
+        self.type = type_
+        self.value = value
+
+    def __repr__(self):
+        return f"BaseDirectiveToken({self.type}, {self.value})"
+
+
+class BaseTokenizer:
+    TOKEN_TYPES = {
+        "KEYWORD": r"\b(BEGIN|END|IF|ELSE|FOR|IN|TRY|EXCEPT|ENDIF|ENDFOR|ENDTRY|DO)\b",
+        "OPERATOR": r"(==|!=|>=|<=|>|<|&&|\|\||!)",
+        "FUNCTION_CALL": r"\b[a-zA-Z_][a-zA-Z0-9_]*\b\((.*?)\)",
+        "LITERAL": r'(\d+|\'.*?\'|".*?")',
+        "IDENTIFIER": r"\b[a-zA-Z_][a-zA-Z0-9_]*\b",
+        "PUNCTUATION": r"(;|,|\(|\))",
+        "WHITESPACE": r"\s+",
+    }
+
+    def __init__(self, script):
+        self.script = script
+        self.tokens = []
+        self.tokenize()
+
+    @property
+    def is_empty(self):
+        return self.tokens == []
+
+    def tokenize(self):
+        position = 0
+        while position < len(self.script):
+            match = None
+            for type_, pattern in self.TOKEN_TYPES.items():
+                regex = re.compile(pattern)
+                match = regex.match(self.script, position)
+                if match:
+                    if type_ != "WHITESPACE":  # Ignore whitespace
+                        token = BaseToken(type_, match.group())
+                        self.tokens.append(token)
+                    position = match.end()  # Move past the matched token
+                    break
+            if not match:  # No match found, unrecognized token
+                raise SyntaxError(f"Unexpected character: {self.script[position]}")
+                # break
+
+    def get_tokens(self):
+        if self.is_empty:
+            try:
+                self.tokenize()
+            except SyntaxError as e:
+                print(e)
+                return []
+        return self.tokens