examples: schedule-extract.py update

Author: pchalasani

examples/basic/chat-azure-client.py

@@ -2,6 +2,17 @@
 Example showing how to use Langroid with Azure OpenAI and Entra ID
 authentication by providing a custom client.
 
+NOTE: this example is ONLY meant for those who are trying to use a custom
+Azure client, as in this scenario:
+https://langroid.github.io/langroid/notes/custom-azure-client/
+This NOT TYPICAL for most users, and should be ignored if you are not using such a
+custom client.
+
+For typical usage of Azure-deployed models with Langroid, see
+the [`test_azure_openai.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_azure_openai.py) and
+[`example/basic/chat.py`](https://github.com/langroid/langroid/blob/main/examples/basic/chat.py)
+
+
 For an async version of this example, see chat-azure-async-client.py.
 
 For more details see here:

examples/basic/chat.py

@@ -71,6 +71,7 @@ def main(
     load_dotenv()
 
     # use the appropriate config instance depending on model name
+    # NOTE: when using Azure, change this to `lm.AzureConfig`
     llm_config = lm.OpenAIGPTConfig(
         chat_model=model or lm.OpenAIChatModel.GPT4o,
         chat_context_length=4096,

examples/basic/python-code-exec-tool.py

@@ -0,0 +1,119 @@
+"""
+Agent that uses a Tool to execute python code.
+
+CAUTION - this is a security risk, as it allows arbitrary code execution.
+This is a bare-bones example. For a real application, you would want to restrict
+the code in various ways, e.g. by using a sandboxed environment, or by restricting
+the modules that can be imported.
+
+Run like this (leave model empty to use default GPT4o)
+
+uv run examples/basic/python-code-exec-tool.py -m gpt4o-mini
+"""
+
+import io
+import contextlib
+from fire import Fire
+from rich.prompt import Prompt
+from langroid.pydantic_v1 import Field
+from langroid.agent.tools.orchestration import ResultTool
+import langroid as lr
+import langroid.language_models as lm
+
+
+def execute_code(code_string):
+    """
+    A minimal function to execute Python code and capture its output.
+
+    Args:
+        code_string: The Python code to execute
+
+    Returns:
+        Tuple of (output, local_variables)
+    """
+    # Create dictionary for local variables
+    local_vars = {}
+
+    # Capture stdout
+    buffer = io.StringIO()
+
+    # Execute code with stdout redirection
+    with contextlib.redirect_stdout(buffer):
+        try:
+            exec(code_string, globals(), local_vars)
+            success = True
+        except Exception as e:
+            print(f"Error: {str(e)}")
+            success = False
+
+    output = buffer.getvalue()
+    return output, local_vars, success
+
+
+class PyCodeTool(lr.ToolMessage):
+    request: str = "py_code_tool"
+    purpose: str = "To execute python <code> and return results"
+
+    code: str = Field(
+        ...,
+        description="""
+            Syntactically valid Python code that can be placed in file to 
+            be run by the Python interpreter. MUST NOT CONTAIN any CODE-BLOCK
+            delimiters like triple-backticks.
+            """,
+    )
+
+    def handle(self):
+        output, local_vars, success = execute_code(self.code)
+        if success:
+            print("Successfully ran code. Results:")
+            print(output)
+            print("Local variables:")
+            print(local_vars)
+        else:
+            print("Failed to run code.")
+        return ResultTool(output=output, local_vars=local_vars, success=success)
+
+
+def main(model: str = ""):
+    llm_config = lm.OpenAIGPTConfig(
+        chat_model=model or lm.OpenAIChatModel.GPT4o,
+    )
+    agent = lr.ChatAgent(
+        lr.ChatAgentConfig(
+            name="Coder",
+            llm=llm_config,
+            # handle LLM non-tool msg
+            handle_llm_no_tool=lambda msg: ResultTool(
+                output=msg.content,
+                success=True,
+            ),
+            system_message=f"""
+            You are an expert python coder. When you get a user's message, 
+            respond as follows:
+            - if you think you need to run Python code,
+                use the TOOL `{PyCodeTool.name()}` to perform the task.
+            - otherwise simply respond to the user's message.
+            """,
+        )
+    )
+    agent.enable_message(PyCodeTool)
+    # task specialized to return ResultTool
+    # set restart to False to maintain conv history across `run` calls
+    task = lr.Task(agent, interactive=False, restart=False)[ResultTool]
+
+    while True:
+        user_input = Prompt.ask("User")
+        if user_input.lower() in ["x", "q"]:
+            break
+        result: ResultTool | None = task.run(user_input)
+        if result is not None:
+            # code was run; do something with the output if any
+            if result.success:
+                print("Output:", result.output)
+            else:
+                print("Code execution failed.")
+
+
+if __name__ == "__main__":
+    Fire(main)

examples/basic/schedule-extract.py

@@ -4,21 +4,25 @@
 Enter vague, unstructured info like:
 
 M-F 8-3pm at home or Tue/Wed 9-1030am at daycare
+
+Run like this -- (omit the -m arg for default gpt-4o-mini LLM)
+
+```bash
+uv run examples/basic/schedule-extract.py -m gpt-4o
 """
 
 import langroid as lr
 import langroid.language_models as lm
-from enum import Enum
 from langroid.agent.tools.orchestration import FinalResultTool
-from typing import List, Dict, Tuple
+from typing import List, Dict, Literal, Tuple
 from langroid.pydantic_v1 import BaseModel, Field
 from rich.prompt import Prompt
 from fire import Fire
 
 
 class Slot(BaseModel):
-    start_time: float = Field(..., description="start time of the slot, e.g. 11:30AM")
-    duration: float = Field(..., description="duration of the slot in MINUTES")
+    start_time: str = Field(..., description="start time of the slot, e.g. 11:30AM")
+    end_time: str = Field(..., description="end time of the slot, e.g. 12:30PM")
     location: str = Field(..., description="location of the slot or UNKNOWN")
 
 
@@ -30,28 +34,19 @@ class DaySchedule(BaseModel):
     slots: List[Slot] = Field(..., description="List of time slots for the day")
 
 
-class Weekday(int, Enum):
-    """
-    A class to represent a weekday.
-    """
-
-    MON = 0
-    TUE = 1
-    WED = 2
-    THU = 3
-    FRI = 4
+Weekday = Literal["Mon", "Tue", "Wed", "Thu", "Fri"]
 
 
 class Availability(BaseModel):
     """
     A class to represent schedule information.
     """
 
-    week_availability: Dict[int, DaySchedule] = Field(
+    week_availability: Dict[Weekday, DaySchedule] = Field(
         ...,
         description="""
         Dictionary mapping weekday to DaySchedule,
-        where 0 = Monday, 1 = Tuesday, ... 4 = Friday
+        where weekday is one of "Mon", "Tue", "Wed", "Thu", "Fri"
         """,
     )
 
@@ -77,17 +72,27 @@ def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
                 cls(
                     availabilities=Availability(
                         week_availability={
-                            Weekday.MON: DaySchedule(
+                            "Mon": DaySchedule(
                                 slots=[
-                                    Slot(start_time=10, duration=360, location="home"),
                                     Slot(
-                                        start_time=15, duration=60, location="daycare"
+                                        start_time="10:00",
+                                        end_time="16:00",
+                                        location="home",
+                                    ),
+                                    Slot(
+                                        start_time="15:00",
+                                        end_time="16:00",
+                                        location="daycare",
                                     ),
                                 ]
                             ),
-                            Weekday.WED: DaySchedule(
+                            "Wed": DaySchedule(
                                 slots=[
-                                    Slot(start_time=10, duration=360, location="home")
+                                    Slot(
+                                        start_time="10:00",
+                                        end_time="16:00",
+                                        location="home",
+                                    )
                                 ]
                             ),
                         }
@@ -110,7 +115,7 @@ def handle(self) -> str:
 
 def make_schedule_task(model: str = ""):
     llm_config = lm.OpenAIGPTConfig(
-        chat_model=model or lm.GeminiModel.GEMINI_2_FLASH_LITE,
+        chat_model=model or lm.OpenAIChatModel.GPT4o_MINI,
     )
     agent = lr.ChatAgent(
         lr.ChatAgentConfig(

examples/docqa/chat.py

@@ -85,6 +85,7 @@ def main(
     config = DocChatAgentConfig(
         llm=llm_config,
         n_query_rephrases=0,
+        full_citations=True,
         hypothetical_answer=False,
         # how many sentences in each segment, for relevance-extraction:
         # increase this if you find that relevance extraction is losing context

examples/docqa/chat_search.py

@@ -16,13 +16,16 @@
     -f: use OpenAI functions api instead of tools
     -m <model_name>:  run with a specific LLM
     (defaults to GPT4-Turbo if blank)
+    -c <crawler_name>: specify a crawler to use for web search. Options are:
+         "trafilatura" (default), "firecrawl"
 
 See here for guide to using local LLMs with Langroid:
 https://langroid.github.io/langroid/tutorials/local-llm-setup/
 """
 
+import typer
 import re
-from typing import List, Any
+from typing import List, Any, Optional
 
 from rich import print
 from rich.prompt import Prompt
@@ -41,6 +44,7 @@
 from langroid.utils.constants import NO_ANSWER
 from langroid.utils.configuration import set_global, Settings
 from fire import Fire
+from langroid.parsing.url_loader import TrafilaturaConfig, FirecrawlConfig
 
 
 class RelevantExtractsTool(ToolMessage):
@@ -85,6 +89,24 @@ def instructions(cls) -> str:
 
 class SearchDocChatAgent(DocChatAgent):
     tried_vecdb: bool = False
+    crawler: Optional[str] = None
+
+    def __init__(self, config: DocChatAgentConfig, crawler: Optional[str] = None):
+        super().__init__(config)
+        self.tried_vecdb = False
+        self.crawler = crawler
+        self.update_crawler_config(crawler)
+
+    def update_crawler_config(self, crawler: Optional[str]):
+        """Updates the crawler config based on the crawler argument."""
+        if crawler == "firecrawl":
+            self.config.crawler_config = FirecrawlConfig()
+        elif crawler == "trafilatura" or crawler is None:
+            self.config.crawler_config = TrafilaturaConfig()
+        else:
+            raise ValueError(
+                f"Unsupported crawler {crawler}. Options are: 'trafilatura', 'firecrawl'"
+            )
 
     def llm_response(
         self,
@@ -127,12 +149,32 @@ def cli():
     Fire(main)
 
 
+app = typer.Typer()
+
+
+@app.command()
 def main(
     debug: bool = False,
     nocache: bool = False,
     model: str = "",
     fn_api: bool = True,
+    crawler: Optional[str] = typer.Option(
+        None,
+        "--crawler",
+        "-c",
+        help="Specify a crawler to use (trafilatura, firecrawl)",
+    ),
 ) -> None:
+    """
+    Main function to run the chatbot.
+
+    Args:
+        debug (bool): Enable debug mode.
+        nocache (bool): Disable caching.
+        model (str): Specify the LLM model to use.
+        fn_api (bool): Use OpenAI functions API instead of tools.
+        crawler (str): Specify the crawler to use for web search.
+    """
 
     set_global(
         Settings(
@@ -169,7 +211,7 @@ def main(
         # "ollama/llama2"
         # "local/localhost:8000/v1"
         # "local/localhost:8000"
-        chat_context_length=2048,  # adjust based on model
+        chat_context_length=8000,  # adjust based on model
     )
 
     config = DocChatAgentConfig(
@@ -188,7 +230,7 @@ def main(
         3. If you are still unable to answer, you can use the `relevant_search_extracts`
            tool/function-call to get some text from a web search. Once you receive the
            text, you can use it to answer my question.
-        4. If you still can't answer, simply say {NO_ANSWER} 
+        5. If you still can't answer, simply say {NO_ANSWER} 
         
         Remember to always FIRST try `relevant_extracts` to see if there are already 
         any relevant docs, before trying web-search with `relevant_search_extracts`.
@@ -204,7 +246,7 @@ def main(
         """,
     )
 
-    agent = SearchDocChatAgent(config)
+    agent = SearchDocChatAgent(config, crawler=crawler)
     agent.enable_message(RelevantExtractsTool)
     agent.enable_message(RelevantSearchExtractsTool)
     collection_name = Prompt.ask(
@@ -225,8 +267,10 @@ def main(
     agent.vecdb.set_collection(collection_name, replace=replace)
 
     task = Task(agent, interactive=False)
-    task.run("Can you help me answer some questions, possibly using web search?")
+    task.run(
+        "Can you help me answer some questions, possibly using web search and crawling?"
+    )
 
 
 if __name__ == "__main__":
-    Fire(main)
+    app()