langroid 0.1.230__py3-none-any.whl → 0.1.233__py3-none-any.whl

langroid/agent/batch.py

@@ -31,6 +31,8 @@ def run_batch_task_gen(
     turns: int = -1,
     message: Optional[str] = None,
     handle_exceptions: bool = False,
+    max_cost: float = 0.0,
+    max_tokens: int = 0,
 ) -> list[U]:
     """
     Generate and run copies of a task async/concurrently one per item in `items` list.
@@ -50,6 +52,8 @@ def run_batch_task_gen(
         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
+        max_cost: float: maximum cost to run the task (default 0.0 for unlimited)
+
 
     Returns:
         list[Any]: list of final results
@@ -62,7 +66,9 @@ def run_batch_task_gen(
             task_i.agent.llm.set_stream(False)
         task_i.agent.config.show_stats = False
 
-        result = await task_i.run_async(input, turns=turns)
+        result = await task_i.run_async(
+            input, turns=turns, max_cost=max_cost, max_tokens=max_tokens
+        )
         return result
 
     async def _do_all(
@@ -120,6 +126,8 @@ def run_batch_tasks(
     sequential: bool = True,
     batch_size: Optional[int] = None,
     turns: int = -1,
+    max_cost: float = 0.0,
+    max_tokens: int = 0,
 ) -> List[U]:
     """
     Run copies of `task` async/concurrently one per item in `items` list.
@@ -137,6 +145,8 @@ def run_batch_tasks(
         batch_size (Optional[int]): The number of tasks to run at a time,
             if None, unbatched
         turns (int): number of turns to run, -1 for infinite
+        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)
 
     Returns:
         list[Any]: list of final results
@@ -151,6 +161,8 @@ def run_batch_tasks(
         batch_size,
         turns,
         message,
+        max_cost=max_cost,
+        max_tokens=max_tokens,
     )
 
 

langroid/agent/special/doc_chat_agent.py

@@ -117,6 +117,7 @@ class DocChatAgentConfig(ChatAgentConfig):
     )
     rerank_diversity: bool = True  # rerank to maximize diversity?
     rerank_periphery: bool = True  # rerank to avoid Lost In the Middle effect?
+    rerank_after_adding_context: bool = True  # rerank after adding context window?
     embed_batch_size: int = 500  # get embedding of at most this many at a time
     cache: bool = True  # cache results
     debug: bool = False
@@ -1110,9 +1111,10 @@ class DocChatAgent(ChatAgent):
         if len(passages) == 0:
             return []
 
-        passages_scores = [(p, 0.0) for p in passages]
-        passages_scores = self.add_context_window(passages_scores)
-        passages = [p for p, _ in passages_scores]
+        if self.config.rerank_after_adding_context:
+            passages_scores = [(p, 0.0) for p in passages]
+            passages_scores = self.add_context_window(passages_scores)
+            passages = [p for p, _ in passages_scores]
         # now passages can potentially have a lot of doc chunks,
         # so we re-rank them using a cross-encoder scoring model,
         # and pick top k where k = config.parsing.n_similar_docs
@@ -1129,6 +1131,11 @@ class DocChatAgent(ChatAgent):
             # (see Lost In the Middle issue).
             passages = self.rerank_to_periphery(passages)
 
+        if not self.config.rerank_after_adding_context:
+            passages_scores = [(p, 0.0) for p in passages]
+            passages_scores = self.add_context_window(passages_scores)
+            passages = [p for p, _ in passages_scores]
+
         return passages
 
     @no_type_check

langroid/agent/task.py

@@ -170,7 +170,8 @@ class Task:
                 agent.set_system_message(system_message)
             if user_message:
                 agent.set_user_message(user_message)
-
+        self.max_cost: float = 0
+        self.max_tokens: int = 0
         self.logger: None | RichFileLogger = None
         self.tsv_logger: None | logging.Logger = None
         self.color_log: bool = False if settings.notebook else True
@@ -375,11 +376,15 @@ class Task:
         msg: Optional[str | ChatDocument] = None,
         turns: int = -1,
         caller: None | Task = None,
+        max_cost: float = 0,
+        max_tokens: int = 0,
     ) -> Optional[ChatDocument]:
         """Synchronous version of `run_async()`.
         See `run_async()` for details."""
         self.task_progress = False
         self.n_stalled_steps = 0
+        self.max_cost = max_cost
+        self.max_tokens = max_tokens
         assert (
             msg is None or isinstance(msg, str) or isinstance(msg, ChatDocument)
         ), f"msg arg in Task.run() must be None, str, or ChatDocument, not {type(msg)}"
@@ -418,6 +423,8 @@ class Task:
         msg: Optional[str | ChatDocument] = None,
         turns: int = -1,
         caller: None | Task = None,
+        max_cost: float = 0,
+        max_tokens: int = 0,
     ) -> Optional[ChatDocument]:
         """
         Loop over `step()` until task is considered done or `turns` is reached.
@@ -434,6 +441,8 @@ class Task:
             turns (int): number of turns to run the task for;
                 default is -1, which means run until task is done.
             caller (Task|None): the calling task, if any
+            max_cost (float): max cost allowed for the task (default 0 -> no limit)
+            max_tokens (int): max tokens allowed for the task (default 0 -> no limit)
 
         Returns:
             Optional[ChatDocument]: valid result of the task.
@@ -444,6 +453,9 @@ class Task:
         # message can be considered to be from the USER
         # (from the POV of this agent's LLM).
         self.task_progress = False
+        self.n_stalled_steps = 0
+        self.max_cost = max_cost
+        self.max_tokens = max_tokens
         if (
             isinstance(msg, ChatDocument)
             and msg.metadata.recipient != ""
@@ -819,6 +831,8 @@ class Task:
                 self.pending_message,
                 turns=actual_turns,
                 caller=self,
+                max_cost=self.max_cost,
+                max_tokens=self.max_tokens,
             )
             result_str = str(ChatDocument.to_LLMMessage(result))
             maybe_tool = len(extract_top_level_json(result_str)) > 0
@@ -894,6 +908,8 @@ class Task:
                 self.pending_message,
                 turns=actual_turns,
                 caller=self,
+                max_cost=self.max_cost,
+                max_tokens=self.max_tokens,
             )
             result_str = str(ChatDocument.to_LLMMessage(result))
             maybe_tool = len(extract_top_level_json(result_str)) > 0
@@ -1051,6 +1067,25 @@ class Task:
             )
             return True
 
+        if self.max_cost > 0 and self.agent.llm is not None:
+            try:
+                if self.agent.llm.tot_tokens_cost()[1] > self.max_cost:
+                    logger.warning(
+                        f"Task {self.name} cost exceeded {self.max_cost}; exiting."
+                    )
+                    return True
+            except Exception:
+                pass
+
+        if self.max_tokens > 0 and self.agent.llm is not None:
+            try:
+                if self.agent.llm.tot_tokens_cost()[0] > self.max_tokens:
+                    logger.warning(
+                        f"Task {self.name} uses > {self.max_tokens} tokens; exiting."
+                    )
+                    return True
+            except Exception:
+                pass
         return (
             # no valid response from any entity/agent in current turn
             result is None

langroid/language_models/base.py

@@ -449,6 +449,18 @@ class LanguageModel(ABC):
             s += f"{model}: {counter}\n"
         return s
 
+    @classmethod
+    def tot_tokens_cost(cls) -> Tuple[int, float]:
+        """
+        Return total tokens used and total cost across all models.
+        """
+        total_tokens = 0
+        total_cost = 0.0
+        for counter in cls.usage_cost_dict.values():
+            total_tokens += counter.total_tokens
+            total_cost += counter.cost
+        return total_tokens, total_cost
+
     def followup_to_standalone(
         self, chat_history: List[Tuple[str, str]], question: str
     ) -> str:

{langroid-0.1.230.dist-info → langroid-0.1.233.dist-info}/METADATA

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

{langroid-0.1.230.dist-info → langroid-0.1.233.dist-info}/RECORD

@@ -1,7 +1,7 @@
 langroid/__init__.py,sha256=qgY-OqzYSWOc6EytQJN9sH2PwDp1UIzP9lXhrYH6aLU,1645
 langroid/agent/__init__.py,sha256=_D8dxnfdr92ch1CIrUkKjrB5HVvsQdn62b1Fb2kBxV8,785
 langroid/agent/base.py,sha256=jyGFmojrFuOy81lUkNsJlR6mLIOY6kOD20P9dhEcEuw,35059
-langroid/agent/batch.py,sha256=T9dgSPThrmIWxQxqDlGwhHa7yw3XIKE_U30bLMRDpNQ,9481
+langroid/agent/batch.py,sha256=SyvUwKetPH_4JKgTnV0bgXS7eD_qfFl9cq06lnill4o,9956
 langroid/agent/callbacks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/callbacks/chainlit.py,sha256=aYuJ8M4VDHr5oymoXL2bpThM7p6P9L45fgJf3MLdkWo,20997
 langroid/agent/chat_agent.py,sha256=X5uVMm9qdw3j-FRf4hbN8k8ByaSdtQCTuU8olKE0sbs,38750
@@ -10,7 +10,7 @@ langroid/agent/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/junk,sha256=LxfuuW7Cijsg0szAzT81OjWWv1PMNI-6w_-DspVIO2s,339
 langroid/agent/openai_assistant.py,sha256=xWSDR4SyMcZhCUkgaM4-mu77rbEDF_xpt7N8m8CkiA4,32962
 langroid/agent/special/__init__.py,sha256=XPE076zD-roskxNBn-A1hnh4AHoMiQN9gk1UDjPaBaU,1201
-langroid/agent/special/doc_chat_agent.py,sha256=Ur57J3xolTaoTgzr2DfP8TY7BViIMNFGNEmmRfHbkL4,53373
+langroid/agent/special/doc_chat_agent.py,sha256=LwWNb_1s5n9rOk9OpOFPuuY1VnVX5DjzQmPwBanKRrM,53763
 langroid/agent/special/lance_doc_chat_agent.py,sha256=USp0U3eTaJzwF_3bdqE7CedSLbaqAi2tm-VzygcyLaA,10175
 langroid/agent/special/lance_rag/__init__.py,sha256=QTbs0IVE2ZgDg8JJy1zN97rUUg4uEPH7SLGctFNumk4,174
 langroid/agent/special/lance_rag/critic_agent.py,sha256=pi_9eMBxEycbWTddtq_yz-mOb2V4SgGm3zfsOH1HU-Q,5775
@@ -32,7 +32,7 @@ langroid/agent/special/sql/utils/populate_metadata.py,sha256=x2OMKfmIBnJESBG3qKt
 langroid/agent/special/sql/utils/system_message.py,sha256=qKLHkvQWRQodTtPLPxr1GSLUYUFASZU8x-ybV67cB68,1885
 langroid/agent/special/sql/utils/tools.py,sha256=6uB2424SLtmapui9ggcEr0ZTiB6_dL1-JRGgN8RK9Js,1332
 langroid/agent/special/table_chat_agent.py,sha256=-Qtqr2FP8VcyYcA-Pzqa9ucSl1-nXudbNsv_qakSSco,9041
-langroid/agent/task.py,sha256=sFncTES0L_O2IpbnHRrNzad0HhtsbHoQ7j3Cjc9eTt0,49711
+langroid/agent/task.py,sha256=jUeKRreYvs_6OlOqVxqfYWQHtvDZViUDtx5wl_b9cyY,51160
 langroid/agent/tool_message.py,sha256=2kPsQUwi3ZzINTUNj10huKnZLjLp5SXmefacTHx8QDc,8304
 langroid/agent/tools/__init__.py,sha256=q-maq3k2BXhPAU99G0H6-j_ozoRvx15I1RFpPVicQIU,304
 langroid/agent/tools/duckduckgo_search_tool.py,sha256=mLGhlgs6pwbYZIwrOs9shfh1dMBVT4DtkR29pYL3cCQ,1900
@@ -59,7 +59,7 @@ langroid/embedding_models/protoc/embeddings_pb2_grpc.py,sha256=9dYQqkW3JPyBpSEje
 langroid/embedding_models/remote_embeds.py,sha256=6_kjXByVbqhY9cGwl9R83ZcYC2km-nGieNNAo1McHaY,5151
 langroid/language_models/__init__.py,sha256=5L9ndEEC8iLJHjDJmYFTnv6-2-3xsxWUMHcugR8IeDs,821
 langroid/language_models/azure_openai.py,sha256=ncRCbKooqLVOY-PWQUIo9C3yTuKEFbAwyngXT_M4P7k,5989
-langroid/language_models/base.py,sha256=Yy_6TP9Qj5CmNtDVQfbcyfytCsvGyow0e1OeqhWGY0A,20638
+langroid/language_models/base.py,sha256=B6dX43ZR65mIvjD95W4RcfpT-WpmiuEcstR3eMrr56Y,21029
 langroid/language_models/config.py,sha256=5UF3DzO1a-Dfsc3vghE0XGq7g9t_xDsRCsuRiU4dgBg,366
 langroid/language_models/openai_assistants.py,sha256=9K-DEAL2aSWHeXj2hwCo2RAlK9_1oCPtqX2u1wISCj8,36
 langroid/language_models/openai_gpt.py,sha256=hR0ibmumIq1TnOGdb80SW0EN6YZn6DZazd5waFHKqGk,49357
@@ -120,7 +120,7 @@ langroid/vector_store/meilisearch.py,sha256=d2huA9P-NoYRuAQ9ZeXJmMKr7ry8u90RUSR2
 langroid/vector_store/momento.py,sha256=9cui31TTrILid2KIzUpBkN2Ey3g_CZWOQVdaFsA4Ors,10045
 langroid/vector_store/qdrant_cloud.py,sha256=3im4Mip0QXLkR6wiqVsjV1QvhSElfxdFSuDKddBDQ-4,188
 langroid/vector_store/qdrantdb.py,sha256=foKRxRv0BBony6S4Vt0Vav9Rn9HMxZvcIh1cE7nosFE,13524
-langroid-0.1.230.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
-langroid-0.1.230.dist-info/METADATA,sha256=SfBCAgSG2P-_1R2NxEFQWLFTAVzdNOGXRhR79EpYkeE,47863
-langroid-0.1.230.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-langroid-0.1.230.dist-info/RECORD,,
+langroid-0.1.233.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
+langroid-0.1.233.dist-info/METADATA,sha256=vxJscoG8qWCn8tfLfCDhVHiomACst_EJcldbP3PNw_M,47863
+langroid-0.1.233.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+langroid-0.1.233.dist-info/RECORD,,