langroid 0.1.100__py3-none-any.whl → 0.1.102__py3-none-any.whl

langroid/parsing/search.py

@@ -7,7 +7,6 @@ See tests for examples: tests/main/test_string_search.py
 """
 
 import difflib
-import re
 from typing import List, Tuple
 
 from nltk.corpus import stopwords
@@ -24,6 +23,7 @@ from .utils import download_nltk_resource
 def find_fuzzy_matches_in_docs(
     query: str,
     docs: List[Document],
+    docs_clean: List[Document],
     k: int,
     words_before: int | None = None,
     words_after: int | None = None,
@@ -49,45 +49,45 @@ def find_fuzzy_matches_in_docs(
         return []
     best_matches = process.extract(
         query,
-        [d.content for d in docs],
+        [d.content for d in docs_clean],
         limit=k,
         scorer=fuzz.partial_ratio,
     )
 
     real_matches = [m for m, score in best_matches if score > 50]
-
-    results = []
-    for match in real_matches:
-        words = match.split()
-        for doc in docs:
-            if match in doc.content:
-                words_in_text = doc.content.split()
-                first_word_idx = next(
-                    (
-                        i
-                        for i, word in enumerate(words_in_text)
-                        if word.startswith(words[0])
-                    ),
-                    -1,
-                )
-                if words_before is None:
-                    words_before = len(words_in_text)
-                if words_after is None:
-                    words_after = len(words_in_text)
-                if first_word_idx != -1:
-                    start_idx = max(0, first_word_idx - words_before)
-                    end_idx = min(
-                        len(words_in_text),
-                        first_word_idx + len(words) + words_after,
-                    )
-                    doc_match = Document(
-                        content=" ".join(words_in_text[start_idx:end_idx]),
-                        metadata=doc.metadata,
-                    )
-                    results.append(doc_match)
+    # find the original docs that corresponding to the matches
+    orig_doc_matches = []
+    for i, m in enumerate(real_matches):
+        for j, doc_clean in enumerate(docs_clean):
+            if m in doc_clean.content:
+                orig_doc_matches.append(docs[j])
                 break
+    if words_after is None and words_before is None:
+        return orig_doc_matches
+
+    contextual_matches = []
+    for match in orig_doc_matches:
+        choice_text = match.content
+        contexts = []
+        while choice_text != "":
+            context, start_pos, end_pos = get_context(
+                query, choice_text, words_before, words_after
+            )
+            if context == "" or end_pos == 0:
+                break
+            contexts.append(context)
+            words = choice_text.split()
+            end_pos = min(end_pos, len(words))
+            choice_text = " ".join(words[end_pos:])
+        if len(contexts) > 0:
+            contextual_matches.append(
+                Document(
+                    content=" ... ".join(contexts),
+                    metadata=match.metadata,
+                )
+            )
 
-    return results
+    return contextual_matches
 
 
 def preprocess_text(text: str) -> str:
@@ -171,7 +171,7 @@ def get_context(
     text: str,
     words_before: int | None = 100,
     words_after: int | None = 100,
-) -> str:
+) -> Tuple[str, int, int]:
     """
     Returns a portion of text containing the best approximate match of the query,
     including b words before and a words after the match.
@@ -185,7 +185,9 @@ def get_context(
     Returns:
     str: A string containing b words before, the match, and a words after
         the best approximate match position of the query in the text. If no
-        match is found, returns "No match found".
+        match is found, returns empty string.
+    int: The start position of the match in the text.
+    int: The end position of the match in the text.
 
     Example:
     >>> get_context("apple", "The quick brown fox jumps over the apple.", 3, 2)
@@ -193,26 +195,29 @@ def get_context(
     """
     if words_after is None and words_before is None:
         # return entire text since we're not asked to return a bounded context
-        return text
+        return text, 0, 0
+
+    # make sure there is a good enough fu
+    if fuzz.partial_ratio(query, text) < 70:
+        return "", 0, 0
 
     sequence_matcher = difflib.SequenceMatcher(None, text, query)
     match = sequence_matcher.find_longest_match(0, len(text), 0, len(query))
 
     if match.size == 0:
-        return "No match found"
-
-    words = re.findall(r"\b\w+\b", text)
-    if words_after is None:
-        words_after = len(words)
-    if words_before is None:
-        words_before = len(words)
-    start_word_pos = len(re.findall(r"\b\w+\b", text[: match.a]))
-    start_pos = max(0, start_word_pos - words_before)
-    end_pos = min(
-        len(words), start_word_pos + words_after + len(re.findall(r"\b\w+\b", query))
-    )
+        return "", 0, 0
+
+    segments = text.split()
+    n_segs = len(segments)
+
+    start_segment_pos = len(text[: match.a].split())
+
+    words_before = words_before or n_segs
+    words_after = words_after or n_segs
+    start_pos = max(0, start_segment_pos - words_before)
+    end_pos = min(len(segments), start_segment_pos + words_after + len(query.split()))
 
-    return " ".join(words[start_pos:end_pos])
+    return " ".join(segments[start_pos:end_pos]), start_pos, end_pos
 
 
 def eliminate_near_duplicates(passages: List[str], threshold: float = 0.8) -> List[str]:

langroid/parsing/utils.py

@@ -165,6 +165,32 @@ def parse_number_range_list(specs: str) -> List[int]:
     return sorted(list(spec_indices))
 
 
+def strip_k(s: str, k: int = 2) -> str:
+    """
+    Strip any leading and trailing whitespaces from the input text beyond length k.
+    This is useful for removing leading/trailing whitespaces from a text while
+    preserving paragraph structure.
+
+    Args:
+        s (str): The input text.
+        k (int): The number of leading and trailing whitespaces to retain.
+
+    Returns:
+        str: The text with leading and trailing whitespaces removed beyond length k.
+    """
+
+    # Count leading and trailing whitespaces
+    leading_count = len(s) - len(s.lstrip())
+    trailing_count = len(s) - len(s.rstrip())
+
+    # Determine how many whitespaces to retain
+    leading_keep = min(leading_count, k)
+    trailing_keep = min(trailing_count, k)
+
+    # Use slicing to get the desired output
+    return s[leading_count - leading_keep : len(s) - (trailing_count - trailing_keep)]
+
+
 def clean_whitespace(text: str) -> str:
     """Remove extra whitespace from the input text, while preserving
     paragraph structure.

langroid/utils/algorithms/graph.py

@@ -0,0 +1,49 @@
+"""
+Graph algos.
+"""
+
+from typing import List, no_type_check
+
+import numpy as np
+
+
+@no_type_check
+def topological_sort(order: np.array) -> List[int]:
+    """
+    Given a directed adjacency matrix, return a topological sort of the nodes.
+    order[i,j] = -1 means there is an edge from i to j.
+    order[i,j] = 0 means there is no edge from i to j.
+    order[i,j] = 1 means there is an edge from j to i.
+
+    Args:
+        order (np.array): The adjacency matrix.
+
+    Returns:
+        List[int]: The topological sort of the nodes.
+
+    """
+    n = order.shape[0]
+
+    # Calculate the in-degrees
+    in_degree = [0] * n
+    for i in range(n):
+        for j in range(n):
+            if order[i, j] == -1:
+                in_degree[j] += 1
+
+    # Initialize the queue with nodes of in-degree 0
+    queue = [i for i in range(n) if in_degree[i] == 0]
+    result = []
+
+    while queue:
+        node = queue.pop(0)
+        result.append(node)
+
+        for i in range(n):
+            if order[node, i] == -1:
+                in_degree[i] -= 1
+                if in_degree[i] == 0:
+                    queue.append(i)
+
+    assert len(result) == n, "Cycle detected"
+    return result

langroid/utils/configuration.py

@@ -1,5 +1,7 @@
+import copy
 import os
-from typing import List
+from contextlib import contextmanager
+from typing import Iterator, List
 
 from dotenv import find_dotenv, load_dotenv
 from pydantic import BaseSettings
@@ -17,6 +19,7 @@ class Settings(BaseSettings):
     gpt3_5: bool = True  # use GPT-3.5?
     nofunc: bool = False  # use model without function_call? (i.e. gpt-4)
     chat_model: str = ""  # language model name, e.g. litellm/ollama/llama2
+    quiet: bool = False  # quiet mode (i.e. suppress all output)?
 
     class Config:
         extra = "forbid"
@@ -55,6 +58,32 @@ def set_global(key_vals: Settings) -> None:
     settings.__dict__.update(key_vals.__dict__)
 
 
+@contextmanager
+def temporary_settings(temp_settings: Settings) -> Iterator[None]:
+    """Temporarily update the global settings and restore them afterward."""
+    original_settings = copy.deepcopy(settings)
+
+    set_global(temp_settings)
+
+    try:
+        yield
+    finally:
+        settings.__dict__.update(original_settings.__dict__)
+
+
+@contextmanager
+def quiet_mode() -> Iterator[None]:
+    """Temporarily set quiet=True in global settings and restore afterward."""
+    original_quiet = settings.quiet
+
+    set_global(Settings(quiet=True))
+
+    try:
+        yield
+    finally:
+        settings.quiet = original_quiet
+
+
 def set_env(settings: BaseSettings) -> None:
     """
     Set environment variables from a BaseSettings instance

langroid/utils/output/printing.py

@@ -1,5 +1,6 @@
 import sys
-from typing import Any, Optional
+from contextlib import contextmanager
+from typing import Any, Iterator, Optional
 
 from rich import print as rprint
 from rich.text import Text
@@ -46,3 +47,32 @@ class PrintColored:
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         print(Colors().RESET)
+
+
+@contextmanager
+def silence_stdout() -> Iterator[None]:
+    """
+    Temporarily silence all output to stdout and from rich.print.
+
+    This context manager redirects all output written to stdout (which includes
+    outputs from the built-in print function and rich.print) to /dev/null on
+    UNIX-like systems or NUL on Windows. Once the context block exits, stdout is
+    restored to its original state.
+
+    Example:
+        with silence_stdout_and_rich():
+            print("This won't be printed")
+            rich.print("This also won't be printed")
+
+    Note:
+        This suppresses both standard print functions and the rich library outputs.
+    """
+    platform_null = "/dev/null" if sys.platform != "win32" else "NUL"
+    original_stdout = sys.stdout
+    fnull = open(platform_null, "w")
+    sys.stdout = fnull
+    try:
+        yield
+    finally:
+        sys.stdout = original_stdout
+        fnull.close()

langroid/utils/pydantic_utils.py

@@ -79,7 +79,9 @@ def flatten_pydantic_model(
         current_model, current_prefix = models_to_process.pop()
 
         for name, field in current_model.__fields__.items():
-            if issubclass(field.outer_type_, BaseModel):
+            if isinstance(field.outer_type_, type) and issubclass(
+                field.outer_type_, BaseModel
+            ):
                 new_prefix = (
                     f"{current_prefix}{name}__" if current_prefix else f"{name}__"
                 )

langroid/vector_store/base.py

@@ -1,12 +1,16 @@
+import copy
 import logging
 from abc import ABC, abstractmethod
-from typing import List, Optional, Sequence, Tuple
+from math import ceil
+from typing import Dict, List, Optional, Sequence, Tuple
 
+import numpy as np
 from pydantic import BaseSettings
 
 from langroid.embedding_models.base import EmbeddingModelsConfig
 from langroid.embedding_models.models import OpenAIEmbeddingsConfig
 from langroid.mytypes import Document
+from langroid.utils.algorithms.graph import topological_sort
 from langroid.utils.configuration import settings
 from langroid.utils.output.printing import print_long_text
 
@@ -130,8 +134,160 @@ class VectorStore(ABC):
         k: int = 1,
         where: Optional[str] = None,
     ) -> List[Tuple[Document, float]]:
+        """
+        Find k most similar texts to the given text, in terms of vector distance metric
+        (e.g., cosine similarity).
+
+        Args:
+            text (str): The text to find similar texts for.
+            k (int, optional): Number of similar texts to retrieve. Defaults to 1.
+            where (Optional[str], optional): Where clause to filter the search.
+
+        Returns:
+            List[Tuple[Document,float]]: List of (Document, score) tuples.
+
+        """
         pass
 
+    def add_context_window(
+        self, docs_scores: List[Tuple[Document, float]], neighbors: int = 0
+    ) -> List[Tuple[Document, float]]:
+        """
+        In each doc's metadata, there may be a window_ids field indicating
+        the ids of the chunks around the current chunk.
+        These window_ids may overlap, so we
+        - gather connected-components of overlapping windows,
+        - split each component into roughly equal parts,
+        - create a new document for each part, preserving metadata,
+
+        We may have stored a longer set of window_ids than we need.
+        We just want `neighbors` on each side of the center of window_ids.
+
+        Args:
+            docs (List[Document]): List of documents to add context window to.
+            scores (List[float]): List of match scores for each document.
+            neighbors (int, optional): Number of neighbors on "each side" of match to
+                retrieve. Defaults to 0.
+                "Each side" here means before and after the match,
+                in the original text.
+
+        Returns:
+            List[Tuple[Document, float]]: List of (Document, score) tuples.
+        """
+        # We return a larger context around each match, i.e.
+        # a window of `neighbors` on each side of the match.
+        docs = [d for d, s in docs_scores]
+        scores = [s for d, s in docs_scores]
+        if neighbors == 0:
+            return docs_scores
+        doc_chunks = [d for d in docs if d.metadata.is_chunk]
+        if len(doc_chunks) == 0:
+            return docs_scores
+        window_ids_list = []
+        id2metadata = {}
+        # id -> highest score of a doc it appears in
+        id2max_score: Dict[int | str, float] = {}
+        for i, d in enumerate(docs):
+            window_ids = d.metadata.window_ids
+            id2metadata.update({id: d.metadata for id in window_ids})
+
+            id2max_score.update(
+                {id: max(id2max_score.get(id, 0), scores[i]) for id in window_ids}
+            )
+            n = len(window_ids)
+            chunk_idx = window_ids.index(d.id())
+            neighbor_ids = window_ids[
+                max(0, chunk_idx - neighbors) : min(n, chunk_idx + neighbors + 1)
+            ]
+            window_ids_list += [neighbor_ids]
+
+        # window_ids could be from different docs,
+        # and they may overlap, so we first remove overlaps
+        window_ids_list = self.remove_overlaps(window_ids_list)
+        final_docs = []
+        final_scores = []
+        for w in window_ids_list:
+            metadata = copy.deepcopy(id2metadata[w[0]])
+            metadata.window_ids = w
+            document = Document(
+                content=" ".join([d.content for d in self.get_documents_by_ids(w)]),
+                metadata=metadata,
+            )
+            # make a fresh id since content is in general different
+            document.metadata.id = document.hash_id(document.content)
+            final_docs += [document]
+            final_scores += [max(id2max_score[id] for id in w)]
+        return list(zip(final_docs, final_scores))
+
+    @staticmethod
+    def remove_overlaps(windows: List[List[str]]) -> List[List[str]]:
+        """
+        Given a collection of windows, where each window is a sequence of ids,
+        identify groups of overlapping windows, and for each overlapping k-group,
+        split the ids into k roughly equal sequences.
+
+        Args:
+            windows (List[int|str]): List of windows, where each window is a
+                sequence of ids.
+
+        Returns:
+            List[int|str]: List of windows, where each window is a sequence of ids,
+                and no two windows overlap.
+        """
+        ids = set(id for w in windows for id in w)
+        # id -> {win -> # pos}
+        id2win2pos: Dict[str, Dict[int, int]] = {id: {} for id in ids}
+
+        for i, w in enumerate(windows):
+            for j, id in enumerate(w):
+                id2win2pos[id][i] = j
+
+        n = len(windows)
+        # relation between windows:
+        order = np.zeros((n, n), dtype=np.int8)
+        for i, w in enumerate(windows):
+            for j, x in enumerate(windows):
+                if i == j:
+                    continue
+                if len(set(w).intersection(x)) == 0:
+                    continue
+                id = list(set(w).intersection(x))[0]  # any common id
+                if id2win2pos[id][i] > id2win2pos[id][j]:
+                    order[i, j] = -1  # win i is before win j
+                else:
+                    order[i, j] = 1  # win i is after win j
+
+        # find groups of windows that overlap, like connected components in a graph
+        groups = [[0]]
+        for i in range(1, n):
+            found = False
+            for g in groups:
+                if any(order[i, j] != 0 for j in g):
+                    g.append(i)
+                    found = True
+                    break
+            if not found:
+                groups.append([i])
+
+        # split each group into roughly equal parts
+        new_windows = []
+        max_window_len = max(len(w) for w in windows)
+        for g in groups:
+            # find total ordering among windows in group based on order matrix
+            # (this is a topological sort)
+            _g = np.array(g)
+            order_matrix = order[_g][:, _g]
+            ordered_window_indices = topological_sort(order_matrix)
+            ordered_window_ids = [windows[i] for i in _g[ordered_window_indices]]
+            flattened = [id for w in ordered_window_ids for id in w]
+            flattened_deduped = list(dict.fromkeys(flattened))
+            # split into k parts where k is the smallest integer such that
+            # each part has length <= max_window_len
+            k = max(1, int(ceil(len(flattened_deduped) / max_window_len)))
+            new_windows += np.array_split(flattened_deduped, k)
+
+        return [w.tolist() for w in new_windows]
+
     @abstractmethod
     def get_all_documents(self) -> List[Document]:
         """

langroid/vector_store/chromadb.py

@@ -109,14 +109,17 @@ class ChromaDB(VectorStore):
         if documents is None:
             return
         contents: List[str] = [document.content for document in documents]
-        metadatas: List[dict[str, Any]] = [
-            document.metadata.dict() for document in documents
-        ]
+        # convert metadatas to dicts so chroma can handle them
+        metadata_dicts: List[dict[str, Any]] = [d.metadata.dict() for d in documents]
+        for m in metadata_dicts:
+            # chroma does not handle non-atomic types in metadata
+            m["window_ids"] = ",".join(m["window_ids"])
+
         ids = [str(d.id()) for d in documents]
         self.collection.add(
             # embedding_models=embedding_models,
             documents=contents,
-            metadatas=metadatas,
+            metadatas=metadata_dicts,
             ids=ids,
         )
 
@@ -145,7 +148,8 @@ class ChromaDB(VectorStore):
             include=["documents", "distances", "metadatas"],
         )
         docs = self._docs_from_results(results)
-        scores = results["distances"][0]
+        # chroma distances are 1 - cosine.
+        scores = [1 - s for s in results["distances"][0]]
         return list(zip(docs, scores))
 
     def _docs_from_results(self, results: Dict[str, Any]) -> List[Document]:
@@ -164,22 +168,11 @@ class ChromaDB(VectorStore):
             for i, c in enumerate(contents):
                 print_long_text("red", "italic red", f"MATCH-{i}", c)
         metadatas = results["metadatas"][0]
+        for m in metadatas:
+            # restore the stringified list of window_ids into the original List[str]
+            m["window_ids"] = m["window_ids"].split(",")
         docs = [
             Document(content=d, metadata=DocMetaData(**m))
             for d, m in zip(contents, metadatas)
         ]
         return docs
-
-
-# Example usage and testing
-# chroma_db = ChromaDB.from_documents(
-#     collection_name="all-my-documents",
-#     documents=["doc1000101", "doc288822"],
-#     metadatas=[{"style": "style1"}, {"style": "style2"}],
-#     ids=["uri9", "uri10"]
-# )
-# results = chroma_db.query(
-#     query_texts=["This is a query document"],
-#     n_results=2
-# )
-# print(results)

langroid/vector_store/meilisearch.py

@@ -263,6 +263,7 @@ class MeiliSearch(VectorStore):
         text: str,
         k: int = 20,
         where: Optional[str] = None,
+        neighbors: int = 0,  # ignored
     ) -> List[Tuple[Document, float]]:
         filter = [] if where is None else where
         if self.config.collection_name is None:

langroid/vector_store/momento.py

@@ -222,6 +222,7 @@ class MomentoVI(VectorStore):
         text: str,
         k: int = 1,
         where: Optional[str] = None,
+        neighbors: int = 0,  # ignored
     ) -> List[Tuple[Document, float]]:
         if self.config.collection_name is None:
             raise ValueError("No collection name set, cannot search")

langroid/vector_store/qdrantdb.py

@@ -244,7 +244,11 @@ class QdrantDB(VectorStore):
             with_vectors=False,
             with_payload=True,
         )
-        docs = [Document(**record.payload) for record in records]  # type: ignore
+        # Note the records may NOT be in the order of the ids,
+        # so we re-order them here.
+        id2payload = {record.id: record.payload for record in records}
+        ordered_payloads = [id2payload[id] for id in _ids]
+        docs = [Document(**payload) for payload in ordered_payloads]  # type: ignore
         return docs
 
     def similar_texts_with_scores(
@@ -252,6 +256,7 @@ class QdrantDB(VectorStore):
         text: str,
         k: int = 1,
         where: Optional[str] = None,
+        neighbors: int = 0,
     ) -> List[Tuple[Document, float]]:
         embedding = self.embedding_fn([text])[0]
         # TODO filter may not work yet
@@ -268,7 +273,7 @@ class QdrantDB(VectorStore):
                 exact=False,  # use Apx NN, not exact NN
             ),
         )
-        scores = [match.score for match in search_result]
+        scores = [match.score for match in search_result if match is not None]
         docs = [
             Document(**(match.payload))  # type: ignore
             for match in search_result
@@ -277,8 +282,9 @@ class QdrantDB(VectorStore):
         if len(docs) == 0:
             logger.warning(f"No matches found for {text}")
             return []
-        if settings.debug:
-            logger.info(f"Found {len(docs)} matches, max score: {max(scores)}")
         doc_score_pairs = list(zip(docs, scores))
+        max_score = max(ds[1] for ds in doc_score_pairs)
+        if settings.debug:
+            logger.info(f"Found {len(doc_score_pairs)} matches, max score: {max_score}")
         self.show_if_debug(doc_score_pairs)
         return doc_score_pairs

{langroid-0.1.100.dist-info → langroid-0.1.102.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.1.100
+Version: 0.1.102
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani