langroid 0.33.12__py3-none-any.whl → 0.34.0__py3-none-any.whl

langroid/agent/batch.py

@@ -1,7 +1,19 @@
 import asyncio
 import copy
 import inspect
-from typing import Any, Callable, Coroutine, Iterable, List, Optional, TypeVar, cast
+import warnings
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    Coroutine,
+    Iterable,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from dotenv import load_dotenv
 
@@ -21,6 +33,238 @@ T = TypeVar("T")
 U = TypeVar("U")
 
 
+class ExceptionHandling(str, Enum):
+    """Enum for exception handling options."""
+
+    RAISE = "raise"
+    RETURN_NONE = "return_none"
+    RETURN_EXCEPTION = "return_exception"
+
+
+def _convert_exception_handling(
+    handle_exceptions: Union[bool, ExceptionHandling]
+) -> ExceptionHandling:
+    """Convert legacy boolean handle_exceptions to ExceptionHandling enum."""
+    if isinstance(handle_exceptions, ExceptionHandling):
+        return handle_exceptions
+
+    if isinstance(handle_exceptions, bool):
+        warnings.warn(
+            "Boolean handle_exceptions is deprecated. "
+            "Use ExceptionHandling enum instead: "
+            "RAISE, RETURN_NONE, or RETURN_EXCEPTION.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return (
+            ExceptionHandling.RETURN_NONE
+            if handle_exceptions
+            else ExceptionHandling.RAISE
+        )
+
+    raise TypeError(
+        "handle_exceptions must be bool or ExceptionHandling, "
+        f"not {type(handle_exceptions)}"
+    )
+
+
+async def _process_batch_async(
+    inputs: Iterable[str | ChatDocument],
+    do_task: Callable[[str | ChatDocument, int], Coroutine[Any, Any, Any]],
+    start_idx: int = 0,
+    stop_on_first_result: bool = False,
+    sequential: bool = False,
+    handle_exceptions: Union[bool, ExceptionHandling] = ExceptionHandling.RAISE,
+    output_map: Callable[[Any], Any] = lambda x: x,
+) -> List[Optional[ChatDocument] | BaseException]:
+    """
+    Unified batch processing logic for both agent methods and tasks.
+
+    Args:
+        inputs: Iterable of inputs to process
+        do_task: Task execution function that takes (input, index) and returns result
+        start_idx: Starting index for the batch
+        stop_on_first_result: Whether to stop after first valid result
+        sequential: Whether to process sequentially
+        handle_exceptions: How to handle exceptions:
+            - RAISE or False: Let exceptions propagate
+            - RETURN_NONE or True: Convert exceptions to None in results
+            - RETURN_EXCEPTION: Include exception objects in results
+            Boolean values are deprecated and will be removed in a future version.
+        output_map: Function to map results to final output format
+    """
+    exception_handling = _convert_exception_handling(handle_exceptions)
+
+    def handle_error(e: BaseException) -> Any:
+        """Handle exceptions based on exception_handling."""
+        match exception_handling:
+            case ExceptionHandling.RAISE:
+                raise e
+            case ExceptionHandling.RETURN_NONE:
+                return None
+            case ExceptionHandling.RETURN_EXCEPTION:
+                return e
+
+    if stop_on_first_result:
+        results: List[Optional[ChatDocument] | BaseException] = []
+        pending: set[asyncio.Task[Any]] = set()
+        # Create task-to-index mapping
+        task_indices: dict[asyncio.Task[Any], int] = {}
+        try:
+            tasks = [
+                asyncio.create_task(do_task(input, i + start_idx))
+                for i, input in enumerate(inputs)
+            ]
+            task_indices = {task: i for i, task in enumerate(tasks)}
+            results = [None] * len(tasks)
+
+            done, pending = await asyncio.wait(
+                tasks, return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # Process completed tasks
+            for task in done:
+                index = task_indices[task]
+                try:
+                    result = await task
+                    results[index] = output_map(result)
+                except BaseException as e:
+                    results[index] = handle_error(e)
+
+            if any(r is not None for r in results):
+                return results
+        finally:
+            for task in pending:
+                task.cancel()
+            try:
+                await asyncio.gather(*pending, return_exceptions=True)
+            except BaseException as e:
+                handle_error(e)
+        return results
+
+    elif sequential:
+        results = []
+        for i, input in enumerate(inputs):
+            try:
+                result = await do_task(input, i + start_idx)
+                results.append(output_map(result))
+            except BaseException as e:
+                results.append(handle_error(e))
+        return results
+
+    # Parallel execution
+    else:
+        try:
+            return_exceptions = exception_handling != ExceptionHandling.RAISE
+            with quiet_mode(), SuppressLoggerWarnings():
+                results_with_exceptions = cast(
+                    list[Optional[ChatDocument | BaseException]],
+                    await asyncio.gather(
+                        *(
+                            do_task(input, i + start_idx)
+                            for i, input in enumerate(inputs)
+                        ),
+                        return_exceptions=return_exceptions,
+                    ),
+                )
+
+                if exception_handling == ExceptionHandling.RETURN_NONE:
+                    results = [
+                        None if isinstance(r, BaseException) else r
+                        for r in results_with_exceptions
+                    ]
+                else:  # ExceptionHandling.RETURN_EXCEPTION
+                    results = results_with_exceptions
+        except BaseException as e:
+            results = [handle_error(e) for _ in inputs]
+
+        return [output_map(r) for r in results]
+
+
+def run_batched_tasks(
+    inputs: List[str | ChatDocument],
+    do_task: Callable[[str | ChatDocument, int], Coroutine[Any, Any, Any]],
+    batch_size: Optional[int],
+    stop_on_first_result: bool,
+    sequential: bool,
+    handle_exceptions: Union[bool, ExceptionHandling],
+    output_map: Callable[[Any], Any],
+    message_template: str,
+    message: Optional[str] = None,
+) -> List[Any]:
+    """
+    Common batch processing logic for both agent methods and tasks.
+
+    Args:
+        inputs: List of inputs to process
+        do_task: Task execution function
+        batch_size: Size of batches, if None process all at once
+        stop_on_first_result: Whether to stop after first valid result
+        sequential: Whether to process sequentially
+        handle_exceptions: How to handle exceptions:
+            - RAISE or False: Let exceptions propagate
+            - RETURN_NONE or True: Convert exceptions to None in results
+            - RETURN_EXCEPTION: Include exception objects in results
+            Boolean values are deprecated and will be removed in a future version.
+        output_map: Function to map results
+        message_template: Template for status message
+        message: Optional override for status message
+    """
+
+    async def run_all_batched_tasks(
+        inputs: List[str | ChatDocument],
+        batch_size: int | None,
+    ) -> List[Any]:
+        """Extra wrap to run asyncio.run one single time and not once per loop
+
+        Args:
+            inputs (List[str  |  ChatDocument]): inputs to process
+            batch_size (int | None): batch size
+
+        Returns:
+            List[Any]: results
+        """
+        results: List[Any] = []
+        if batch_size is None:
+            msg = message or message_template.format(total=len(inputs))
+            with status(msg), SuppressLoggerWarnings():
+                results = await _process_batch_async(
+                    inputs,
+                    do_task,
+                    stop_on_first_result=stop_on_first_result,
+                    sequential=sequential,
+                    handle_exceptions=handle_exceptions,
+                    output_map=output_map,
+                )
+        else:
+            batches = batched(inputs, batch_size)
+            for batch in batches:
+                start_idx = len(results)
+                complete_str = f", {start_idx} complete" if start_idx > 0 else ""
+                msg = (
+                    message or message_template.format(total=len(inputs)) + complete_str
+                )
+
+                if stop_on_first_result and any(r is not None for r in results):
+                    results.extend([None] * len(batch))
+                else:
+                    with status(msg), SuppressLoggerWarnings():
+                        results.extend(
+                            await _process_batch_async(
+                                batch,
+                                do_task,
+                                start_idx=start_idx,
+                                stop_on_first_result=stop_on_first_result,
+                                sequential=sequential,
+                                handle_exceptions=handle_exceptions,
+                                output_map=output_map,
+                            )
+                        )
+        return results
+
+    return asyncio.run(run_all_batched_tasks(inputs, batch_size))
+
+
 def run_batch_task_gen(
     gen_task: Callable[[int], Task],
     items: list[T],
@@ -31,7 +275,7 @@ def run_batch_task_gen(
     batch_size: Optional[int] = None,
     turns: int = -1,
     message: Optional[str] = None,
-    handle_exceptions: bool = False,
+    handle_exceptions: Union[bool, ExceptionHandling] = ExceptionHandling.RAISE,
     max_cost: float = 0.0,
     max_tokens: int = 0,
 ) -> list[Optional[U]]:
@@ -58,7 +302,11 @@ def run_batch_task_gen(
             if None, unbatched
         turns (int): number of turns to run, -1 for infinite
         message (Optional[str]): optionally overrides the console status messages
-        handle_exceptions: bool: Whether to replace exceptions with outputs of None
+        handle_exceptions: How to handle exceptions:
+            - RAISE or False: Let exceptions propagate
+            - RETURN_NONE or True: Convert exceptions to None in results
+            - RETURN_EXCEPTION: Include exception objects in results
+            Boolean values are deprecated and will be removed in a future version.
         max_cost: float: maximum cost to run the task (default 0.0 for unlimited)
         max_tokens: int: maximum token usage (in and out) (default 0 for unlimited)
 
@@ -72,113 +320,33 @@ def run_batch_task_gen(
     async def _do_task(
         input: str | ChatDocument,
         i: int,
-        return_idx: Optional[int] = None,
     ) -> BaseException | Optional[ChatDocument] | tuple[int, Optional[ChatDocument]]:
         task_i = gen_task(i)
         if task_i.agent.llm is not None:
             task_i.agent.llm.set_stream(False)
         task_i.agent.config.show_stats = False
+
         try:
             result = await task_i.run_async(
                 input, turns=turns, max_cost=max_cost, max_tokens=max_tokens
             )
-            if return_idx is not None:
-                return return_idx, result
-            else:
-                return result
         except asyncio.CancelledError as e:
             task_i.kill()
-            if handle_exceptions:
-                return e
-            else:
-                raise e
-        except BaseException as e:
-            if handle_exceptions:
-                return e
-            else:
-                raise e
-
-    async def _do_all(
-        inputs: Iterable[str | ChatDocument], start_idx: int = 0
-    ) -> list[Optional[U]]:
-        results: list[Optional[ChatDocument]] = []
-        if stop_on_first_result:
-            outputs: list[Optional[U]] = [None] * len(list(inputs))
-            tasks = set(
-                asyncio.create_task(_do_task(input, i + start_idx, return_idx=i))
-                for i, input in enumerate(inputs)
-            )
-            while tasks:
-                try:
-                    done, tasks = await asyncio.wait(
-                        tasks, return_when=asyncio.FIRST_COMPLETED
-                    )
-                    for task in done:
-                        idx_result = task.result()
-                        if not isinstance(idx_result, tuple):
-                            continue
-                        index, output = idx_result
-                        outputs[index] = output_map(output)
-
-                    if any(r is not None for r in outputs):
-                        return outputs
-                finally:
-                    # Cancel all remaining tasks
-                    for task in tasks:
-                        task.cancel()
-                    # Wait for cancellations to complete
-                    try:
-                        await asyncio.gather(*tasks, return_exceptions=True)
-                    except BaseException as e:
-                        if not handle_exceptions:
-                            raise e
-            return outputs
-        elif sequential:
-            for i, input in enumerate(inputs):
-                result: Optional[ChatDocument] | BaseException = await _do_task(
-                    input, i + start_idx
-                )  # type: ignore
-
-                if isinstance(result, BaseException):
-                    result = None
-
-                results.append(result)
-        else:
-            results_with_exceptions = cast(
-                list[Optional[ChatDocument | BaseException]],
-                await asyncio.gather(
-                    *(_do_task(input, i + start_idx) for i, input in enumerate(inputs)),
-                ),
-            )
-
-            results = [
-                r if not isinstance(r, BaseException) else None
-                for r in results_with_exceptions
-            ]
-
-        return list(map(output_map, results))
-
-    results: List[Optional[U]] = []
-    if batch_size is None:
-        msg = message or f"[bold green]Running {len(items)} tasks:"
-
-        with status(msg), SuppressLoggerWarnings():
-            results = asyncio.run(_do_all(inputs))
-    else:
-        batches = batched(inputs, batch_size)
-
-        for batch in batches:
-            start_idx = len(results)
-            complete_str = f", {start_idx} complete" if start_idx > 0 else ""
-            msg = message or f"[bold green]Running {len(items)} tasks{complete_str}:"
-
-            if stop_on_first_result and any(r is not None for r in results):
-                results.extend([None] * len(batch))
-            else:
-                with status(msg), SuppressLoggerWarnings():
-                    results.extend(asyncio.run(_do_all(batch, start_idx=start_idx)))
-
-    return results
+            # exception will be handled by the caller
+            raise e
+        return result
+
+    return run_batched_tasks(
+        inputs=inputs,
+        do_task=_do_task,
+        batch_size=batch_size,
+        stop_on_first_result=stop_on_first_result,
+        sequential=sequential,
+        handle_exceptions=handle_exceptions,
+        output_map=output_map,
+        message_template="[bold green]Running {total} tasks:",
+        message=message,
+    )
 
 
 def run_batch_tasks(
@@ -242,6 +410,8 @@ def run_batch_agent_method(
     output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
     sequential: bool = True,
     stop_on_first_result: bool = False,
+    handle_exceptions: Union[bool, ExceptionHandling] = ExceptionHandling.RAISE,
+    batch_size: Optional[int] = None,
 ) -> List[Any]:
     """
     Run the `method` on copies of `agent`, async/concurrently one per
@@ -265,6 +435,14 @@ def run_batch_agent_method(
             to final result
         sequential (bool): whether to run sequentially
             (e.g. some APIs such as ooba don't support concurrent requests)
+        stop_on_first_result (bool): whether to stop after the first valid
+        handle_exceptions: How to handle exceptions:
+            - RAISE or False: Let exceptions propagate
+            - RETURN_NONE or True: Convert exceptions to None in results
+            - RETURN_EXCEPTION: Include exception objects in results
+            Boolean values are deprecated and will be removed in a future version.
+        batch_size (Optional[int]): The number of items to process in each batch.
+            If None, process all items at once.
     Returns:
         List[Any]: list of final results
     """
@@ -288,43 +466,18 @@ def run_batch_agent_method(
         if method_i is None:
             raise ValueError(f"Agent {agent_name} has no method {method_name}")
         result = await method_i(input)
-        return output_map(result)
+        return result
 
-    async def _do_all() -> List[Any]:
-        if stop_on_first_result:
-            tasks = [
-                asyncio.create_task(_do_task(input, i))
-                for i, input in enumerate(inputs)
-            ]
-            results = [None] * len(tasks)
-            try:
-                done, pending = await asyncio.wait(
-                    tasks, return_when=asyncio.FIRST_COMPLETED
-                )
-                for task in done:
-                    index = tasks.index(task)
-                    results[index] = await task
-            finally:
-                for task in pending:
-                    task.cancel()
-                await asyncio.gather(*pending, return_exceptions=True)
-            return results
-        elif sequential:
-            results = []
-            for i, input in enumerate(inputs):
-                result = await _do_task(input, i)
-                results.append(result)
-            return results
-        with quiet_mode(), SuppressLoggerWarnings():
-            return await asyncio.gather(
-                *(_do_task(input, i) for i, input in enumerate(inputs))
-            )
-
-    n = len(items)
-    with status(f"[bold green]Running {n} copies of {agent_name}..."):
-        results = asyncio.run(_do_all())
-
-    return results
+    return run_batched_tasks(
+        inputs=inputs,
+        do_task=_do_task,
+        batch_size=batch_size,
+        stop_on_first_result=stop_on_first_result,
+        sequential=sequential,
+        handle_exceptions=handle_exceptions,
+        output_map=output_map,
+        message_template=f"[bold green]Running {{total}} copies of {agent_name}...",
+    )
 
 
 def llm_response_batch(
@@ -334,6 +487,7 @@ def llm_response_batch(
     output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
     sequential: bool = True,
     stop_on_first_result: bool = False,
+    batch_size: Optional[int] = None,
 ) -> List[Any]:
     return run_batch_agent_method(
         agent,
@@ -343,6 +497,7 @@ def llm_response_batch(
         output_map=output_map,
         sequential=sequential,
         stop_on_first_result=stop_on_first_result,
+        batch_size=batch_size,
     )
 
 
@@ -353,6 +508,7 @@ def agent_response_batch(
     output_map: Callable[[ChatDocument | None], Any] = lambda x: x,
     sequential: bool = True,
     stop_on_first_result: bool = False,
+    batch_size: Optional[int] = None,
 ) -> List[Any]:
     return run_batch_agent_method(
         agent,
@@ -362,6 +518,7 @@ def agent_response_batch(
         output_map=output_map,
         sequential=sequential,
         stop_on_first_result=stop_on_first_result,
+        batch_size=batch_size,
     )
 
 

langroid/agent/chat_agent.py

@@ -1845,14 +1845,16 @@ class ChatAgent(Agent):
         self.update_last_message(message, role=Role.USER)
         return answer_doc
 
-    def llm_response_forget(self, message: str) -> ChatDocument:
+    def llm_response_forget(
+        self, message: Optional[str | ChatDocument] = None
+    ) -> ChatDocument:
         """
         LLM Response to single message, and restore message_history.
         In effect a "one-off" message & response that leaves agent
         message history state intact.
 
         Args:
-            message (str): user message
+            message (str|ChatDocument): message to respond to.
 
         Returns:
             A Document object with the response.
@@ -1879,7 +1881,9 @@ class ChatAgent(Agent):
 
         return response
 
-    async def llm_response_forget_async(self, message: str) -> ChatDocument:
+    async def llm_response_forget_async(
+        self, message: Optional[str | ChatDocument] = None
+    ) -> ChatDocument:
         """
         Async version of `llm_response_forget`. See there for details.
         """

langroid/agent/special/doc_chat_agent.py

@@ -1,3 +1,4 @@
+# # langroid/agent/special/doc_chat_agent.py
 """
 Agent that supports asking queries about a set of documents, using
 retrieval-augmented generation (RAG).
@@ -16,14 +17,14 @@ pip install "langroid[hf-embeddings]"
 import logging
 from collections import OrderedDict
 from functools import cache
-from typing import Any, Dict, List, Optional, Set, Tuple, no_type_check
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, no_type_check
 
 import nest_asyncio
 import numpy as np
 import pandas as pd
 from rich.prompt import Prompt
 
-from langroid.agent.batch import run_batch_tasks
+from langroid.agent.batch import run_batch_agent_method, run_batch_tasks
 from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
 from langroid.agent.chat_document import ChatDocMetaData, ChatDocument
 from langroid.agent.special.relevance_extractor_agent import (
@@ -81,6 +82,8 @@ DEFAULT_DOC_CHAT_SYSTEM_MESSAGE = """
 You are a helpful assistant, helping me understand a collection of documents.
 """
 
+CHUNK_ENRICHMENT_DELIMITER = "<##-##-##>"
+
 has_sentence_transformers = False
 try:
     from sentence_transformers import SentenceTransformer  # noqa: F401
@@ -102,6 +105,12 @@ oai_embed_config = OpenAIEmbeddingsConfig(
 )
 
 
+class ChunkEnrichmentAgentConfig(ChatAgentConfig):
+    batch_size: int = 50
+    delimiter: str = CHUNK_ENRICHMENT_DELIMITER
+    enrichment_prompt_fn: Callable[[str], str] = lambda x: x
+
+
 class DocChatAgentConfig(ChatAgentConfig):
     system_message: str = DEFAULT_DOC_CHAT_SYSTEM_MESSAGE
     user_message: str = DEFAULT_DOC_CHAT_INSTRUCTIONS
@@ -126,6 +135,12 @@ class DocChatAgentConfig(ChatAgentConfig):
     # https://arxiv.org/pdf/2212.10496.pdf
     # It is False by default; its benefits depends on the context.
     hypothetical_answer: bool = False
+    # Optional config for chunk enrichment agent, e.g. to enrich
+    # chunks with hypothetical questions, or keywords to increase
+    # the "semantic surface area" of the chunks, which may help
+    # improve retrieval.
+    chunk_enrichment_config: Optional[ChunkEnrichmentAgentConfig] = None
+
     n_query_rephrases: int = 0
     n_neighbor_chunks: int = 0  # how many neighbors on either side of match to retrieve
     n_fuzzy_neighbor_words: int = 100  # num neighbor words to retrieve for fuzzy match
@@ -404,6 +419,8 @@ class DocChatAgent(ChatAgent):
                 d.metadata.is_chunk = True
         if self.vecdb is None:
             raise ValueError("VecDB not set")
+        if self.config.chunk_enrichment_config is not None:
+            docs = self.enrich_chunks(docs)
 
         # If any additional fields need to be added to content,
         # add them as key=value pairs for all docs, before batching.
@@ -860,6 +877,72 @@ class DocChatAgent(ChatAgent):
                 ).content
         return answer
 
+    def enrich_chunks(self, docs: List[Document]) -> List[Document]:
+        """
+        Enrich chunks using Agent configured with self.config.chunk_enrichment_config.
+
+        We assume that the system message of the agent is set in such a way
+        that when we run
+        ```
+        prompt = self.config.chunk_enrichment_config.enrichment_prompt_fn(text)
+        result = await agent.llm_response_forget_async(prompt)
+        ```
+
+        then `result.content` will contain the augmentation to the text.
+
+        Args:
+            docs: List of document chunks to enrich
+
+        Returns:
+            List[Document]: Documents (chunks) enriched with additional text,
+                separated by a delimiter.
+        """
+        if self.config.chunk_enrichment_config is None:
+            return docs
+        enrichment_config = self.config.chunk_enrichment_config
+        agent = ChatAgent(enrichment_config)
+        if agent.llm is None:
+            raise ValueError("LLM not set")
+
+        with status("[cyan]Augmenting chunks..."):
+            # Process chunks in parallel using run_batch_agent_method
+            questions_batch = run_batch_agent_method(
+                agent=agent,
+                method=agent.llm_response_forget_async,
+                items=docs,
+                input_map=lambda doc: (
+                    enrichment_config.enrichment_prompt_fn(doc.content)
+                ),
+                output_map=lambda response: response.content if response else "",
+                sequential=False,
+                batch_size=enrichment_config.batch_size,
+            )
+
+            # Combine original content with generated questions
+            augmented_docs = []
+            for doc, enrichment in zip(docs, questions_batch):
+                if not enrichment:
+                    augmented_docs.append(doc)
+                    continue
+
+                # Combine original content with questions in a structured way
+                combined_content = f"""
+                {doc.content}
+                
+                {enrichment_config.delimiter}
+                {enrichment}
+                """.strip()
+
+                new_doc = doc.copy(
+                    update={
+                        "content": combined_content,
+                        "metadata": doc.metadata.copy(update={"has_enrichment": True}),
+                    }
+                )
+                augmented_docs.append(new_doc)
+
+            return augmented_docs
+
     def llm_rephrase_query(self, query: str) -> List[str]:
         if self.llm is None:
             raise ValueError("LLM not set")
@@ -1143,20 +1226,22 @@ class DocChatAgent(ChatAgent):
         id2_rank_semantic = {d.id(): i for i, (d, _) in enumerate(docs_and_scores)}
         id2doc = {d.id(): d for d, _ in docs_and_scores}
         # make sure we get unique docs
-        passages = [id2doc[id] for id, _ in id2_rank_semantic.items()]
+        passages = [id2doc[id] for id in id2_rank_semantic.keys()]
 
         id2_rank_bm25 = {}
         if self.config.use_bm25_search:
             # TODO: Add score threshold in config
             docs_scores = self.get_similar_chunks_bm25(query, retrieval_multiple)
+            id2doc.update({d.id(): d for d, _ in docs_scores})
             if self.config.cross_encoder_reranking_model == "":
                 # only if we're not re-ranking with a cross-encoder,
                 # we collect these ranks for Reciprocal Rank Fusion down below.
                 docs_scores = sorted(docs_scores, key=lambda x: x[1], reverse=True)
                 id2_rank_bm25 = {d.id(): i for i, (d, _) in enumerate(docs_scores)}
-                id2doc.update({d.id(): d for d, _ in docs_scores})
             else:
                 passages += [d for (d, _) in docs_scores]
+                # eliminate duplicate ids
+                passages = [id2doc[id] for id in id2doc.keys()]
 
         id2_rank_fuzzy = {}
         if self.config.use_fuzzy_match:
@@ -1174,6 +1259,8 @@ class DocChatAgent(ChatAgent):
                 id2doc.update({d.id(): d for d, _ in fuzzy_match_doc_scores})
             else:
                 passages += [d for (d, _) in fuzzy_match_doc_scores]
+                # eliminate duplicate ids
+                passages = [id2doc[id] for id in id2doc.keys()]
 
         if (
             self.config.cross_encoder_reranking_model == ""
@@ -1301,7 +1388,6 @@ class DocChatAgent(ChatAgent):
         if self.config.n_query_rephrases > 0:
             rephrases = self.llm_rephrase_query(query)
             proxies += rephrases
-
         passages = self.get_relevant_chunks(query, proxies)  # no LLM involved
 
         if len(passages) == 0:
@@ -1315,6 +1401,29 @@ class DocChatAgent(ChatAgent):
 
         return query, extracts
 
+    def remove_chunk_enrichments(self, passages: List[Document]) -> List[Document]:
+        """Remove any enrichments (like hypothetical questions, or keywords)
+        from documents.
+        Only cleans if enrichment was enabled in config.
+
+        Args:
+            passages: List of documents to clean
+
+        Returns:
+            List of documents with only original content
+        """
+        if self.config.chunk_enrichment_config is None:
+            return passages
+        delimiter = self.config.chunk_enrichment_config.delimiter
+        return [
+            (
+                doc.copy(update={"content": doc.content.split(delimiter)[0].strip()})
+                if doc.content and getattr(doc.metadata, "has_enrichment", False)
+                else doc
+            )
+            for doc in passages
+        ]
+
     def get_verbatim_extracts(
         self,
         query: str,
@@ -1330,6 +1439,8 @@ class DocChatAgent(ChatAgent):
         Returns:
             List[Document]: list of Documents containing extracts and metadata.
         """
+        passages = self.remove_chunk_enrichments(passages)
+
         agent_cfg = self.config.relevance_extractor_config
         if agent_cfg is None:
             # no relevance extraction: simply return passages

langroid/mytypes.py

@@ -75,6 +75,17 @@ class Document(BaseModel):
     def id(self) -> str:
         return self.metadata.id
 
+    @staticmethod
+    def from_string(
+        content: str,
+        source: str = "context",
+        is_chunk: bool = True,
+    ) -> "Document":
+        return Document(
+            content=content,
+            metadata=DocMetaData(source=source, is_chunk=is_chunk),
+        )
+
     def __str__(self) -> str:
         return dedent(
             f"""

langroid/parsing/spider.py

@@ -7,9 +7,9 @@ try:
     from pydispatch import dispatcher
     from scrapy import signals
     from scrapy.crawler import CrawlerRunner
-    from scrapy.http import Response
-    from scrapy.linkextractors import LinkExtractor
-    from scrapy.spiders import CrawlSpider, Rule
+    from scrapy.http.response.text import TextResponse
+    from scrapy.linkextractors.lxmlhtml import LxmlLinkExtractor
+    from scrapy.spiders import CrawlSpider, Rule  # type: ignore
     from twisted.internet import defer, reactor
 except ImportError:
     raise LangroidImportError("scrapy", "scrapy")
@@ -21,7 +21,7 @@ class DomainSpecificSpider(CrawlSpider):  # type: ignore
 
     custom_settings = {"DEPTH_LIMIT": 1, "CLOSESPIDER_ITEMCOUNT": 20}
 
-    rules = (Rule(LinkExtractor(), callback="parse_item", follow=True),)
+    rules = (Rule(LxmlLinkExtractor(), callback="parse_item", follow=True),)
 
     def __init__(self, start_url: str, k: int = 20, *args, **kwargs):  # type: ignore
         """Initialize the spider with start_url and k.
@@ -36,13 +36,13 @@ class DomainSpecificSpider(CrawlSpider):  # type: ignore
         self.k = k
         self.visited_urls: Set[str] = set()
 
-    def parse_item(self, response: Response):  # type: ignore
+    def parse_item(self, response: TextResponse):  # type: ignore
         """Extracts URLs that are within the same domain.
 
         Args:
             response: The scrapy response object.
         """
-        for link in LinkExtractor(allow_domains=self.allowed_domains).extract_links(
+        for link in LxmlLinkExtractor(allow_domains=self.allowed_domains).extract_links(
             response
         ):
             if len(self.visited_urls) < self.k:

{langroid-0.33.12.dist-info → langroid-0.34.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langroid
-Version: 0.33.12
+Version: 0.34.0
 Summary: Harness LLMs with Multi-Agent Programming
 Author-email: Prasad Chalasani <pchalasani@gmail.com>
 License: MIT

{langroid-0.33.12.dist-info → langroid-0.34.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
 langroid/__init__.py,sha256=z_fCOLQJPOw3LLRPBlFB5-2HyCjpPgQa4m4iY5Fvb8Y,1800
 langroid/exceptions.py,sha256=gp6ku4ZLdXXCUQIwUNVFojJNGTzGnkevi2PLvG7HOhc,2555
-langroid/mytypes.py,sha256=ptAFxEAtiwmIfUnGisNotTe8wT9LKBf22lOfPgZoQIY,2368
+langroid/mytypes.py,sha256=h1eMq1ZwTLVezObPfCseWNWbEOzP7mAKu2XoS63W1cM,2647
 langroid/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/__init__.py,sha256=ll0Cubd2DZ-fsCMl7e10hf9ZjFGKzphfBco396IKITY,786
 langroid/agent/base.py,sha256=oThlrYygKDu1-bKjAfygldJ511gMKT8Z0qCrD52DdDM,77834
-langroid/agent/batch.py,sha256=qK3ph6VNj_1sOhfXCZY4r6gh035DglDKU751p8BU0tY,14665
-langroid/agent/chat_agent.py,sha256=cxamUgqQkr6_W3mqCPz3L7rJnXIkD4hemR7X7uhlBvI,82095
+langroid/agent/batch.py,sha256=vi1r5i1-vN80WfqHDSwjEym_KfGsqPGUtwktmiK1nuk,20635
+langroid/agent/chat_agent.py,sha256=A-7Iiiw7jsoJNlWerljM29BidkiIbjPOQIkGZpZHmt0,82210
 langroid/agent/chat_document.py,sha256=xPUMGzR83rn4iAEXIw2jy5LQ6YJ6Y0TiZ78XRQeDnJQ,17778
 langroid/agent/openai_assistant.py,sha256=JkAcs02bIrgPNVvUWVR06VCthc5-ulla2QMBzux_q6o,34340
 langroid/agent/task.py,sha256=XrXUbSoiFasvpIsZPn_cBpdWaTCKljJPRimtLMrSZrs,90347
@@ -14,7 +14,7 @@ langroid/agent/xml_tool_message.py,sha256=6SshYZJKIfi4mkE-gIoSwjkEYekQ8GwcSiCv7a
 langroid/agent/callbacks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/callbacks/chainlit.py,sha256=RH8qUXaZE5o2WQz3WJQ1SdFtASGlxWCA6_HYz_3meDQ,20822
 langroid/agent/special/__init__.py,sha256=gik_Xtm_zV7U9s30Mn8UX3Gyuy4jTjQe9zjiE3HWmEo,1273
-langroid/agent/special/doc_chat_agent.py,sha256=zw2MvdCWRPH93d73PKh27KFiQ8sUCFPxAfLDdkxvdZQ,59301
+langroid/agent/special/doc_chat_agent.py,sha256=tI16jVavTSOen9OUoRTl5heDTeTBhWsxW17XU9ZcEko,63563
 langroid/agent/special/lance_doc_chat_agent.py,sha256=s8xoRs0gGaFtDYFUSIRchsgDVbS5Q3C2b2mr3V1Fd-Q,10419
 langroid/agent/special/lance_tools.py,sha256=qS8x4wi8mrqfbYV2ztFzrcxyhHQ0ZWOc-zkYiH7awj0,2105
 langroid/agent/special/relevance_extractor_agent.py,sha256=zIx8GUdVo1aGW6ASla0NPQjYYIpmriK_TYMijqAx3F8,4796
@@ -85,7 +85,7 @@ langroid/parsing/parser.py,sha256=bTG5TO2CEwGdLf9979j9_dFntKX5FloGF8vhts6ObU0,11
 langroid/parsing/repo_loader.py,sha256=3GjvPJS6Vf5L6gV2zOU8s-Tf1oq_fZm-IB_RL_7CTsY,29373
 langroid/parsing/routing.py,sha256=-FcnlqldzL4ZoxuDwXjQPNHgBe9F9-F4R6q7b_z9CvI,1232
 langroid/parsing/search.py,sha256=0i_r0ESb5HEQfagA2g7_uMQyxYPADWVbdcN9ixZhS4E,8992
-langroid/parsing/spider.py,sha256=Y6y7b86Y2k770LdhxgjVlImBxuuy1V9n8-XQ3QPaG5s,3199
+langroid/parsing/spider.py,sha256=hAVM6wxh1pQ0EN4tI5wMBtAjIk0T-xnpi-ZUzWybhos,3258
 langroid/parsing/table_loader.py,sha256=qNM4obT_0Y4tjrxNBCNUYjKQ9oETCZ7FbolKBTcz-GM,3410
 langroid/parsing/url_loader.py,sha256=JK48KktLRDBfjrt4nsUfy92M6yGdEeicAqOum2MdULM,4656
 langroid/parsing/urls.py,sha256=XjpaV5onG7gKQ5iQeFTzHSw5P08Aqw0g-rMUu61lR6s,7988
@@ -121,7 +121,7 @@ langroid/vector_store/lancedb.py,sha256=b3_vWkTjG8mweZ7ZNlUD-NjmQP_rLBZfyKWcxt2v
 langroid/vector_store/meilisearch.py,sha256=6frB7GFWeWmeKzRfLZIvzRjllniZ1cYj3HmhHQICXLs,11663
 langroid/vector_store/momento.py,sha256=UNHGT6jXuQtqY9f6MdqGU14bVnS0zHgIJUa30ULpUJo,10474
 langroid/vector_store/qdrantdb.py,sha256=HRLCt-FG8y4718omwpFaQZnWeYxPj0XCwS4tjokI1sU,18116
-langroid-0.33.12.dist-info/METADATA,sha256=tUh6elP7kcHfPS2dUXBE-gZ6vLoBfFzYsmg7nr2oCrg,59016
-langroid-0.33.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langroid-0.33.12.dist-info/licenses/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
-langroid-0.33.12.dist-info/RECORD,,
+langroid-0.34.0.dist-info/METADATA,sha256=fo7ULfjnWFED6Cag8aUFjOaPqEatQKBXEz-Z_rFyHnk,59015
+langroid-0.34.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langroid-0.34.0.dist-info/licenses/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
+langroid-0.34.0.dist-info/RECORD,,