llms.txt - main - githubmirror/langroid

llms.txt
128 638 строк4 MB
Prasad Chalasani
Support short-lived / rotating credentials via api_key_provider (Vertex AI, Azure AD) (#1081)
27 июл 2026, 19:08
27 июл 2026, 19:088c1a945
This file is a merged representation of a subset of the codebase, containing files not matching ignore patterns, combined into a single document by Repomix.

<file_summary>
This section contains a summary of this file.

<purpose>
This file contains a packed representation of a subset of the repository's contents that is considered the most important context.
It is designed to be easily consumable by AI systems for analysis, code review,
or other automated processes.
</purpose>

<file_format>
The content is organized as follows:
1. This summary section
2. Repository information
3. Directory structure
4. Repository files (if enabled)
5. Multiple file entries, each consisting of:
  - File path as an attribute
  - Full contents of the file
</file_format>

<usage_guidelines>
- This file should be treated as read-only. Any changes should be made to the
  original repository files, not this packed version.
- When processing this file, use the file path to distinguish
  between different files in the repository.
- Be aware that this file may contain sensitive information. Handle it with
  the same level of security as you would the original repository.
- Pay special attention to the Repository Description. These contain important context and guidelines specific to this project.
</usage_guidelines>

<notes>
- Some files may have been excluded based on .gitignore rules and Repomix's configuration
- Binary files are not included in this packed representation. Please refer to the Repository Structure section for a complete list of file paths, including binary files
- Files matching these patterns are excluded: .*/**, !.github, !.github/**, data/**, **/data/**, **/*.csv, logs/**, build/**, dist/**, *.egg-info/**, __pycache__/**, **/*.pyc, .pytest_cache/**, .mypy_cache/**, .ruff_cache/**, venv/**, env/**, .env, **/*.json, llms.txt, llms-compressed.txt, llms-128k.txt, llms-no-tests.txt, llms-no-tests-compressed.txt, llms-no-tests-no-examples.txt, llms-no-tests-no-examples-compressed.txt, file-list.txt, file-list-updated.txt, examples/data/**, examples/docqa/docs/**, examples/logs/**, tests/cache/**, tests/logs/**, **/*.pkl, **/*.pickle, **/*.db, **/*.sqlite, **/*.log, **/node_modules/**, **/*.min.js, **/*.map, coverage/**, htmlcov/**, .coverage, *.orig, *.tmp, *.bak, *.swp, *.swo, **/docker-compose*.yml, visual_log.sh, **/*_converted.md, **/page_*.md, **/page-*.md, tests/main/dummy-pages/**, tests/main/data/**/*.txt, **/*.pb2.py, **/*.pb2_grpc.py
- Files matching patterns in .gitignore are excluded
- Files matching default ignore patterns are excluded
- Files are sorted by Git change count (files with more changes are at the bottom)
</notes>

</file_summary>

<user_provided_header>
Langroid Repository Export for LLM Analysis
</user_provided_header>

<directory_structure>
ai-instructions/
  claude-repomix-instructions.md
ai-notes/
  handler-parameter-analysis-notes.md
  Langroid-repo-docs.md
  repomix-plan.md
docs/
  assets/
    demos/
      dependency_chatbot.gif
      lease-extractor-demo.gif
    DepChatbot.png
    langroid_neo4j_logos.png
    langroid-card-cobalt-blue-1200-630.png
    langroid-card-cobalt-blue.png
    langroid-card-lambda-ossem-rust-1200-630.png
    langroid-card-ossem-rust-1200x630.png
    langroid-card-ossem-rust.png
    langroid-lance.png
    nullify-logo.jpg
    nullify-logo.png
    orange-logo-lambda-563.png
    orange-logo-lambda.png
    orange-logo.png
    pure-lambda-non-circular.png
  blog/
    posts/
      figures/
        confusion.png
        langroid-arch.png
        malade-arch.png
        omop-ground-truth.png
        omop-results.png
        variance-histogram.png
      chat-completion.md
      langroid-architecture.md
      langroid-intro.md
      langroid-knowledge-graph.md
      langroid-lancedb.md
      local-llm-formatting.md
      local-llm.md
      malade.md
      multi-agent-debate.md
      test.md
    .authors.yml
    index.md
  demos/
    targeting/
      audience-targeting.md
      targeting.gif
      targeting.png
  examples/
    agent-tree.md
    chat-tree.png
    guide.md
  javascripts/
    mathjax.js
  logos/
    pure-lambda-non-circular.png
  notes/
    async-streaming.md
    azure-openai-models.md
    batch-processing.md
    chunking.md
    code-injection-protection.md
    crawl4ai.md
    custom-azure-client.md
    enriching-for-retrieval.md
    file-input.md
    gemini.md
    glhf-chat.md
    handle-llm-no-tool.md
    html-logger.md
    knowledge-graphs.md
    langdb.md
    large-tool-results.md
    litellm-proxy.md
    llama-cpp-embeddings.md
    llm-none-content.md
    llm-pdf-parser.md
    marker-pdf.md
    markitdown.md
    mcp-tools.md
    message-routing.md
    openai-client-caching.md
    openai-http-client.md
    overview.md
    pgvector.md
    pinecone.md
    portkey.md
    pydantic-v2-migration.md
    qdrant-resource-cleanup.md
    quiet-mode.md
    reasoning-content.md
    rotating-api-keys.md
    seltz_search.md
    structured-output.md
    task-termination.md
    task-tool.md
    tavily_search.md
    tool-message-handler.md
    twitter-search.md
    url_loader.md
    weaviate.md
    xml-tools.md
  overrides/
    partials/
      comments.html
  quick-start/
    chat-agent-docs.md
    chat-agent-fn.png
    chat-agent-tool.md
    chat-agent-tool.png
    chat-agent.md
    chat-docs.png
    chat.png
    img.png
    index.md
    llm-interaction.md
    multi-agent-task-delegation.md
    setup.md
    three-agent-chat-num-router.md
    three-agent-chat-num.md
    three-agent-num.png
    three-agent-router-func.png
    three-agent-router.png
    try-llm.png
    two-agent-chat-num.md
    two-agent-num.png
    two-agent.png
  screenshots/
    planner-workflow-html-logs.png
  stylesheets/
    extra.css
  tutorials/
    langroid-tour-beethoven.png
    langroid-tour-teacher.png
    langroid-tour.md
    llm-usage-options.md
    local-llm-setup.md
    non-openai-llms.md
    postgresql-agent.md
    supported-models.md
  auto_docstring.py
  FAQ.md
  index.md
examples/
  basic/
    multi-agent-search-critic/
      __init__.py
      assistant_agent.py
      critic_agent.py
      main.py
      search_agent.py
      tools.py
    multi-agent-search-critic-no-orch/
      __init__.py
      assistant_agent.py
      critic_agent.py
      main.py
      search_agent.py
      tools.py
    __init__.py
    1-agent-3-tools-address-user.py
    1-agent-3-tools.py
    1d-screen-click.py
    2-agent-tools.py
    autocorrect.py
    chat-2-agent-discuss.py
    chat-azure-async-client.py
    chat-azure-client.py
    chat-local-numerical.py
    chat-local.py
    chat-minimax.py
    chat-persist.py
    chat-search-assistant-local.py
    chat-search-assistant.py
    chat-search-seltz.py
    chat-search.py
    chat-tool-function.py
    chat-tree-structured-simple.py
    chat-tree-structured.py
    chat-tree.py
    chat.py
    completion.py
    concurrent-tasks.py
    done_sequences_example.py
    drug-outcomes.py
    fn-call-local-numerical.py
    fn-call-local-simple.py
    intent-classifier.py
    multi-agent-medical.py
    multi-agent-return-result.py
    multi-agent-round-table.py
    multi-agent-triage.py
    oai-asst-chat.py
    oai-code-chat.py
    plan-subtasks.py
    planner-workflow-simple.py
    planner-workflow-spawn.py
    planner-workflow.py
    python-code-exec-tool.py
    schedule-extract.py
    text-to-structured.py
    tool-custom-handler.py
    tool-extract-short-example.py
    xml_tool.py
  chainlit/
    non-callback/
      chat-doc-qa-no-callback.py
      chat-no-callback.py
      chat-search-no-callback.py
      chat-stream.py
      chat-tool-no-callback.py
      README.md
    public/
      favicon
      langroid-card.png
      logo_dark.png
      logo_light.png
    books.txt
    chainlit.md
    chat-doc-qa.py
    chat-search-assistant-local.py
    chat-search-assistant.py
    chat-search-rag.py
    chat-search.py
    chat-tool.py
    chat-transcript.py
    chat-tree-chainlit.py
    chat-with-agent.py
    chat-with-task.py
    cypher_message.py
    dependency_chatbot.py
    extract-then-chat.py
    multi-agent-nested-tool.py
    multi-agent.py
    multi-extract-3.py
    multi-extract.py
    README.md
    simplest.py
    test-step-nesting.py
  data-qa/
    sql-chat/
      __init__.py
      sql_chat.py
      utils.py
    __init__.py
    table_chat.py
  docqa/
    streamlit-app/
      app.py
      README.md
      requirements.txt
      utils.py
    __init__.py
    books.txt
    chat_multi_extract.py
    chat_search.py
    chat-local.py
    chat-multi-extract-3.py
    chat-multi-extract-local.py
    chat-qa-summarize.py
    chat-search-filter.py
    chat.py
    crawl4ai_examples.py
    doc-aware-chat.py
    doc-aware-compose-2.py
    doc-aware-guide-2.py
    doc-based-troubleshooting.py
    doc-chat-2.py
    doc-chat-multi-llm.py
    doc-chat-simple.py
    doc-chunk-enrich.py
    extract-then-chat.py
    filter-multi-doc-auto.py
    filter-multi-doc-manual.py
    filter-multi-doc-query-plan.py
    lance-rag-gh-issues.py
    lance-rag-movies.py
    langroid-lancedb-rag-movies.ipynb
    lease.txt
    oai-multi-extract.py
    oai-retrieval-2.py
    oai-retrieval-assistant.py
    rag-concurrent.py
    rag-local-simple.py
  extract/
    capitals.py
    extract.py
    job_listing.txt
    kaggle_text.py
    lease.html
    lease.txt
    least-truncated.txt
    pdf-json-flex.py
    pdf-json-no-parse.py
    pdf-json.py
    README.md
    um-financial-report.pdf
  kg-chat/
    chat-arangodb-igvf.py
    chat-arangodb.py
    chat-neo4j.py
    csv-chat.py
    cypher_message.py
    dependency_chatbot.py
    DependencyChatbot.ipynb
    movies.cypher
    README.md
    text-kg-triplets.py
    text-kg.py
  langdb/
    langdb_chat_agent_docs.py
    langdb_chat_agent_tool.py
    langdb_custom_headers.py
    README.md
    requirements.txt
  mcp/
    any-mcp.py
    biomcp.py
    chainlit-mcp.py
    claude-code-mcp-single.py
    claude-code-mcp.py
    exa-web-search.py
    gitmcp.py
    mcp-fetch.py
    mcp-file-system.py
    memory.py
    openmemory.py
    playwright-mcp.py
    puppeteer-mcp.py
    pyodide_code_executor.py
    xquik-twitter-search.py
  multi-agent-debate/
    chainlit_utils.py
    config.py
    generation_config_models.py
    main_chainlit.py
    main.py
    models.py
    README.md
    system_messages.py
    utils.py
  portkey/
    portkey_advanced_features.py
    portkey_basic_chat.py
    portkey_multi_provider.py
    README.md
    requirements.txt
  privacy/
    annotate.py
    annotate2.py
    privacy_agent.py
    privacy_annotator.py
  quick-start/
    chat-agent-docs.py
    chat-agent-tool.py
    chat-agent.py
    quick-start.ipynb
    three-agent-chat-num-router.py
    three-agent-chat-num.py
    three-agent-chat.py
    try-llm.py
    two-agent-chat-num.py
    two-agent-chat.py
  reasoning/
    agent-reasoning.py
  summarize/
    summ-batch.py
    summ.py
  __init__.py
  langroid_quick_examples.ipynb
  Langroid_quick_start.ipynb
  Langroid_QuickStart_OpenAI_Assistants_API.ipynb
  README.md
issues/
  pydantic-v2-migration/
    examples-errors.md
    migration-checking-log.md
    pr-pydantic-v2-fixes.md
    PYDANTIC_V2_MIGRATION_TASK_SPECIFICATION.md
    pydantic-migration-checking-instructions.md
    PYRANTIC-V2-MIGRATION-PLAN.md
  20251010-concurrent-rag-status.md
  20251010-concurrent-rag.md
  20251011-cross-encoder-race-bug.md
  20251011-pr-926-description.md
  20251107-fix-mcp-dectorator.md
  20251123-new-model-support-gpt51-gemini30.md
  20260601-qdrant-ci-local-container.md
  898-implementation.md
  html-logger-implementation.md
  html-logger.md
  issue-919-llamacpp-embeddings.md
  llm-client-caching-phase1-summary.md
  llm-client-caching-phase2-summary.md
  llm-client-caching-test-summary.md
  llm-client-caching.md
  pr-882-cached-tokens-improvements.md
  pr-openai-client-caching.md
  pr-qdrant-lock-fix.md
  qdrant-lock-issue-spec-changes.md
langroid/
  agent/
    callbacks/
      __init__.py
      chainlit.py
    special/
      arangodb/
        __init__.py
        aql_validator.py
        arangodb_agent.py
        system_messages.py
        tools.py
        utils.py
      lance_rag/
        __init__.py
        critic_agent.py
        lance_rag_task.py
        query_planner_agent.py
      neo4j/
        __init__.py
        csv_kg_chat.py
        cypher_validator.py
        neo4j_chat_agent.py
        system_messages.py
        tools.py
      sql/
        utils/
          __init__.py
          description_extractors.py
          populate_metadata.py
          system_message.py
          tools.py
        __init__.py
        sql_chat_agent.py
      __init__.py
      doc_chat_agent.py
      lance_doc_chat_agent.py
      lance_tools.py
      relevance_extractor_agent.py
      retriever_agent.py
      table_chat_agent.py
    tools/
      mcp/
        __init__.py
        decorators.py
        fastmcp_client.py
      __init__.py
      duckduckgo_search_tool.py
      exa_search_tool.py
      file_tools.py
      google_search_tool.py
      metaphor_search_tool.py
      orchestration.py
      recipient_tool.py
      retrieval_tool.py
      rewind_tool.py
      segment_extract_tool.py
      seltz_search_tool.py
      task_tool.py
      tavily_search_tool.py
    __init__.py
    base.py
    batch.py
    chat_agent.py
    chat_document.py
    done_sequence_parser.py
    openai_assistant.py
    task.py
    tool_message.py
    xml_tool_message.py
  cachedb/
    __init__.py
    base.py
    redis_cachedb.py
  embedding_models/
    protoc/
      __init__.py
      embeddings_pb2_grpc.py
      embeddings_pb2.py
      embeddings_pb2.pyi
      embeddings.proto
    __init__.py
    base.py
    models.py
    remote_embeds.py
  language_models/
    prompt_formatter/
      __init__.py
      base.py
      hf_formatter.py
      llama2_formatter.py
    __init__.py
    azure_openai.py
    base.py
    client_cache.py
    config.py
    mock_lm.py
    model_info.py
    openai_gpt.py
    provider_params.py
    utils.py
  parsing/
    __init__.py
    agent_chats.py
    code_parser.py
    document_parser.py
    file_attachment.py
    md_parser.py
    para_sentence_split.py
    parse_json.py
    parser.py
    pdf_utils.py
    repo_loader.py
    routing.py
    search.py
    spider.py
    table_loader.py
    url_loader.py
    urls.py
    utils.py
    web_search.py
  prompts/
    __init__.py
    dialog.py
    prompts_config.py
    templates.py
  pydantic_v1/
    __init__.py
    main.py
  utils/
    algorithms/
      __init__.py
      graph.py
    output/
      __init__.py
      citations.py
      printing.py
      status.py
    __init__.py
    configuration.py
    constants.py
    git_utils.py
    globals.py
    html_logger.py
    logging.py
    object_registry.py
    pandas_utils.py
    pydantic_utils.py
    system.py
    types.py
  vector_store/
    __init__.py
    base.py
    chromadb.py
    lancedb.py
    meilisearch.py
    pineconedb.py
    postgres.py
    qdrantdb.py
    weaviatedb.py
  __init__.py
  exceptions.py
  mytypes.py
  py.typed
plugins/
  langroid/
    skills/
      add-pattern/
        SKILL.md
      patterns/
        agent-handler-validation-with-state.md
        agent-tool-handler-with-state.md
        done-sequences-specific-tool.md
        mcp-tool-integration.md
        quiet-mode.md
        run-batch-tasks.md
        SKILL.md
        task-return-tool.md
public/
  favicon
  langroid-card.png
  logo_dark.png
  logo_light.png
release-notes/
  v0-56-0-task-tool.md
  v0-56-11-openai-client-caching.md
  v0-56-12-cached-tokens-support.md
  v0-56-13-done-sequences-parent-chain-fixes.md
  v0-56-15-response-sequence-tracking.md
  v0-56-2-table-chat-fix.md
  v0-56-4-handler-params.md
  v0-56-6-doc-chat-refactor.md
  v0-56-7-doc-chat-deprecation-fix.md
  v0-56-8-task-tool-spawn-example.md
  v0-56-9-rrf-crossencoder-fixes.md
  v0-58-0-crawl4ai-integration.md
  v0.57.0-html-logger.md
scripts/
  fix-pydantic-imports.sh
tests/
  extras/
    sql/
      test_automatic_context_extraction.py
    test_csv_kg_chat.py
    test_doc_chat_agent_llamacpp.py
    test_docx_parser_extra.py
    test_fastembed_embeddings.py
    test_gemini_embeddings.py
    test_hf_embeddings.py
    test_hf_vector_stores.py
    test_llamacpp_embedding_formats.py
    test_llamacpp_embeddings.py
    test_marker_pdf_parser.py
    test_pyarango.py
  main/
    mcp/
      weather-server-python/
        pyproject.toml
        README.md
        weather.py
    sql_chat/
      test_sql_chat_agent.py
      test_sql_chat_security.py
    test_agent.py
    test_arango_aql_validation.py
    test_arangodb_chat_agent.py
    test_arangodb.py
    test_async_handlers.py
    test_azure_openai.py
    test_batch_tasks_typed.py
    test_batch.py
    test_callbacks.py
    test_chat_agent_async.py
    test_chat_agent.py
    test_closest_string.py
    test_code_parser.py
    test_concurrent_doc_chat_qdrant.py
    test_concurrent_rag_simple.py
    test_csv_kg_validation.py
    test_dataframe_docs.py
    test_doc_chat_agent.py
    test_doc_chat_relevance.py
    test_docx_parser.py
    test_done_sequence_parser.py
    test_done_sequences_dsl.py
    test_done_sequences.py
    test_embeddings.py
    test_file_attachment.py
    test_file_tools.py
    test_git_utils.py
    test_global_settings.py
    test_global_state.py
    test_handle_message_security.py
    test_html_logger.py
    test_json.py
    test_lance_doc_chat_agent.py
    test_llm_async.py
    test_llm_pdf_parser.py
    test_llm_response.py
    test_llm.py
    test_markitdown_parser.py
    test_mcp_tools.py
    test_md_parser.py
    test_minimax_provider.py
    test_msg_routing.py
    test_multi_agent_complex_async.py
    test_multi_agent_complex.py
    test_multi_agent.py
    test_mytypes.py
    test_neo4j_chat_agent.py
    test_neo4j_cypher_validation.py
    test_object_registry.py
    test_openai_assistant_async.py
    test_openai_assistant.py
    test_openai_gpt_client_cache.py
    test_openai_gpt_none_content.py
    test_openai_http_client_simple.py
    test_openai_http_client.py
    test_openai_params_subclass.py
    test_pandas_utils.py
    test_parser.py
    test_parsing_citations.py
    test_pdf_parser.py
    test_pdf_utils.py
    test_prep_llm_message.py
    test_pydantic_utils.py
    test_qdrant_lock_release.py
    test_quiet_mode.py
    test_recipient_tool_async.py
    test_recipient_tool.py
    test_redis_cache.py
    test_relevance_extractor.py
    test_repo_chunking.py
    test_repo_loader_symlinks.py
    test_repo_loader.py
    test_retriever_agent.py
    test_retry_utils.py
    test_rich_file_logger.py
    test_rotating_api_keys.py
    test_seltz_search.py
    test_split_inline_reasoning.py
    test_stateful_tool.py
    test_stateless_tool_messages.py
    test_string_search.py
    test_structured_output.py
    test_system_utils.py
    test_table_chat_agent.py
    test_task_inf_loop.py
    test_task_lineage_rewind.py
    test_task_optional_logger.py
    test_task_run_polymorphic.py
    test_task_tool.py
    test_task.py
    test_token_usage.py
    test_tool_handler_async.py
    test_tool_handler.py
    test_tool_messages_async.py
    test_tool_messages_azure.py
    test_tool_messages.py
    test_tool_orchestration.py
    test_tool_origin_taint.py
    test_tool_taint_laundering.py
    test_url_loader.py
    test_vector_stores.py
    test_web_search_tools.py
    test_xml_tool_message.py
  __init__.py
  conftest.py
  README.md
  test_pdf_parser_extra.py
  utils.py
.blackignore
.coveragerc
.env-template
.gitignore
.pre-commit-config.yaml
bump_version.sh
chainlit.md
CLAUDE.md
CODE_OF_CONDUCT.md
CONTRIBUTING.md
Dockerfile
LICENSE
Makefile
mkdocs.yml
PR_954_REVIEW.md
PR_REVIEW_975.md
pyproject.toml
pytest.ini
README.md
SECURITY.md
setup.cfg
</directory_structure>

<files>
This section contains the contents of the repository's files.

<file path="ai-instructions/claude-repomix-instructions.md">
# AI Instructions for Setting Up Repomix

## Task Overview
Set up [repomix](https://github.com/yamadashy/repomix) to generate LLM-friendly repository exports. This creates text files that can be uploaded to AI models for code analysis.

## Steps to Complete

### 1. Install Repomix
```bash
npm install -g repomix
```

### 2. Create repomix.config.json
Create a configuration file in the repository root with:
- **Include patterns**: Source code files (*.py, *.js, *.md, *.yaml, *.yml, *.toml)
- **Exclude patterns**: Data directories, logs, node_modules, JSON files, generated files
- **Security check**: Enable to prevent sensitive data inclusion

### 3. Configure Include/Exclude Patterns
- Include only source code directories and documentation
- Exclude data/, logs/, build artifacts, dependencies
- Add `llms*.txt` to exclusions to prevent recursive inclusion

### 4. Test Configuration (Optional)
```bash
# Generate file list only for inspection
repomix --no-files -o file-list.txt
```
This allows you to review which files will be included before generating the full output.

### 5. Generate Output Versions

Use the Makefile targets to generate repomix files:

```bash
# Generate all variants (recommended)
make repomix-all

# Or generate specific versions:
make repomix                      # llms.txt and llms-compressed.txt (includes tests)
make repomix-no-tests             # llms-no-tests.txt and llms-no-tests-compressed.txt
make repomix-no-tests-no-examples # llms-no-tests-no-examples.txt and compressed version
```

All commands use `git ls-files` to ensure only git-tracked files are included.

### 6. Verify Results
- Check file sizes and token counts in repomix output
- Ensure no sensitive data is included
- Confirm only relevant source files are packaged

## Expected Outcome
Six text files optimized for different LLM contexts:
- `llms.txt`: Full version with tests and examples (870K tokens)
- `llms-compressed.txt`: Compressed version with tests and examples (513K tokens)
- `llms-no-tests.txt`: Full version without tests (677K tokens)
- `llms-no-tests-compressed.txt`: Compressed version without tests (433K tokens)
- `llms-no-tests-no-examples.txt`: Core library code only (no tests/examples)
- `llms-no-tests-no-examples-compressed.txt`: Compressed core library code (285K tokens)

The files contain only git-tracked source code with proper exclusions for clean, focused LLM consumption.
</file>

<file path="ai-notes/handler-parameter-analysis-notes.md">
# Handler Parameter Analysis Notes

## Overview

This document summarizes the investigation into how Langroid analyzes handler method parameters in `langroid/agent/base.py`, specifically focusing on the `_analyze_handler_params` method and its role in creating handler wrappers.

## Key Methods and Call Chain

### Call Chain
1. `_get_tool_list()` - Registers tool messages and their handlers
2. `_create_handler_wrapper()` - Creates wrapper functions for handlers
3. `_analyze_handler_params()` - Analyzes handler method signatures

## How _analyze_handler_params Works

The `_analyze_handler_params` method (lines 253-313 in agent/base.py) analyzes a handler method's signature to identify:
- Whether it has type annotations
- Which parameter is the agent parameter
- Which parameter is the chat_doc parameter

### Analysis Process (Updated Implementation)
1. **Type Annotation Check**: First checks if parameters have type annotations
   - **Direct Class Checking** (NEW): For simple class annotations like `Agent` or `ChatAgent`:
     - Uses `inspect.isclass(param.annotation) and issubclass(param.annotation, Agent)`
     - This works because Python stores the actual class object in the annotation
   - **Direct Identity Check** (NEW): For ChatDocument:
     - Uses `param.annotation is ChatDocument` for exact match
   - **String-based Fallback**: For complex type hints like `Optional[Agent]`:
     - Falls back to checking if "Agent" is in the string representation
     - Necessary because complex generic types aren't simple class objects

2. **Fallback to Parameter Names**: If no annotations found
   - Looks for parameter named `agent`
   - Looks for parameter named `chat_doc`

### Key Insight: Type Annotations Are Objects
The crucial realization is that Python's type annotation system stores actual class references when possible:
- `def handler(agent: Agent):` → `param.annotation` contains the actual `Agent` class object
- `def handler(agent: Optional[Agent]):` → `param.annotation` contains a complex type object that requires string inspection
- This allows direct `issubclass()` checks for simple annotations, making the analysis more accurate and robust

## How _create_handler_wrapper Works

Based on the analysis from `_analyze_handler_params`, the wrapper creates different function signatures:
- No parameters → `wrapper(obj)`
- Both agent and chat_doc → `wrapper(obj, chat_doc)` with correct parameter order
- Only agent → `wrapper(obj)` passing agent internally
- Only chat_doc → `wrapper(obj, chat_doc)`

## Why Direct Type Checking Works (Clarification)

Initially, we believed runtime type checking wasn't feasible because we confused two different concepts:

### The Misconception
We thought we needed runtime values to check parameter types, but this was incorrect. The confusion arose from:
1. Thinking we needed actual parameter values to determine their types
2. Not realizing that type annotations are stored as Python objects in the function signature

### The Reality: Static Analysis of Type Annotations
1. **Type annotations are available at definition time**: When Python parses `def handler(agent: Agent):`, it stores the `Agent` class object in the function's signature
2. **No runtime values needed**: We're checking the type annotations themselves, not the runtime values
3. **Direct class comparison is possible**: For simple type hints, `param.annotation` contains the actual class object, allowing `issubclass()` checks

### Why This Approach Works
1. **Setup Time Analysis**: We analyze the handler signature when tools are registered, using the stored annotation objects
2. **Direct Type Checking**: For simple annotations like `Agent`, we can use `issubclass(param.annotation, Agent)`
3. **Fallback for Complex Types**: For generic types like `Optional[Agent]`, we fall back to string matching
4. **Performance**: Still analyzes once at setup, no runtime overhead

## Current Design Benefits
- Analyzes handler signatures once at setup time
- Creates wrappers with exact signatures needed
- No runtime ambiguity about parameter arrangement
- Clear error messages if handler signatures don't match expectations

## Implementation Changes Summary

### Recent Updates to _analyze_handler_params
The method was enhanced to support direct type checking of handler parameters:

1. **Direct Class Checking for Agent Types**:
   ```python
   if inspect.isclass(param.annotation) and issubclass(param.annotation, Agent):
   ```
   - Checks if the annotation is a direct class reference to Agent or its subclasses
   - More accurate than string matching alone

2. **Direct Identity Check for ChatDocument**:
   ```python
   if param.annotation is ChatDocument:
   ```
   - Uses identity comparison for exact ChatDocument type matching

3. **Improved Parameter Extraction**:
   - Changed from `[p for p in params if p.name != "self"]` to `params[1:]`
   - More reliable for removing the 'self' parameter

4. **Fallback Strategy**:
   - Still uses string matching for complex type hints like `Optional[Agent]`
   - Maintains backward compatibility while improving accuracy

## Related PR
This investigation was prompted by PR #861 "MCP updates" which made changes to how `FastMCPServer` forwards image context and resources, and added optional persistence for MCP server connections. The handler parameter analysis improvements were made to support more robust type checking for MCP tool handlers.
</file>

<file path="ai-notes/Langroid-repo-docs.md">
# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Commands

### Development
- Install core dependencies: `pip install -e .`
- Install dev dependencies: `pip install -e ".[dev]"`
- Install specific feature groups:
  - Document chat features: `pip install -e ".[doc-chat]"`
  - Database features: `pip install -e ".[db]"`
  - HuggingFace embeddings: `pip install -e ".[hf-embeddings]"`
  - All features: `pip install -e ".[all]"`
- Run linting and type checking: `make check`
- Format code: `make lint`

### Testing
- Run all tests: `pytest tests/`
- Run specific test: `pytest tests/main/test_file.py::test_function`
- Run tests with coverage: `pytest --cov=langroid tests/`
- Run only main tests: `make tests` (uses `pytest tests/main`)

### Linting and Type Checking
- Lint code: `make check` (runs black, ruff check, mypy)
- Format only: `make lint` (runs black and ruff fix)
- Type check only: `make type-check`
- Always use `make check` to run lints + mypy before trying to commit changes

### Version and Release Management
- Bump version: `./bump_version.sh [patch|minor|major]`
- Or use make commands:
  - `make all-patch` - Bump patch version, build, push, release
  - `make all-minor` - Bump minor version, build, push, release
  - `make all-major` - Bump major version, build, push, release

## Architecture

Langroid is a framework for building LLM-powered agents that can use tools and collaborate with each other.

### Core Components:

1. **Agents** (`langroid/agent/`):
   - `chat_agent.py` - Base ChatAgent that can converse and use tools
   - `task.py` - Handles execution flow for agents
   - `special/` - Domain-specific agents (doc chat, table chat, SQL chat, etc.)
   - `openai_assistant.py` - Integration with OpenAI Assistant API

2. **Tools** (`langroid/agent/tools/`):
   - Tool system for agents to interact with external systems
   - `tool_message.py` - Protocol for tool messages
   - Various search tools (Google, DuckDuckGo, Tavily, Exa, etc.)

3. **Language Models** (`langroid/language_models/`):
   - Abstract interfaces for different LLM providers
   - Implementations for OpenAI, Azure, local models, etc.
   - Support for hundreds of LLMs via LiteLLM

4. **Vector Stores** (`langroid/vector_store/`):
   - Abstract interface and implementations for different vector databases
   - Includes support for Qdrant, Chroma, LanceDB, Pinecone, PGVector, Weaviate

5. **Document Processing** (`langroid/parsing/`):
   - Parse and process documents from various formats
   - Chunk text for embedding and retrieval
   - Support for PDF, DOCX, images, and more

6. **Embedding Models** (`langroid/embedding_models/`):
   - Abstract interface for embedding generation
   - Support for OpenAI, HuggingFace, and custom embeddings

### Key Multi-Agent Patterns:

- **Task Delegation**: Agents can delegate tasks to other agents through hierarchical task structures
- **Message Passing**: Agents communicate by transforming and passing messages
- **Collaboration**: Multiple agents can work together on complex tasks

### Key Security Features:

- The `full_eval` flag in both `TableChatAgentConfig` and `VectorStoreConfig` controls code injection protection
- Defaults to `False` for security, set to `True` only in trusted environments

## Documentation

- Main documentation is in the `docs/` directory
- Examples in the `examples/` directory demonstrate usage patterns
- Quick start examples available in `examples/quick-start/`

## MCP (Model Context Protocol) Tools Integration

Langroid provides comprehensive support for MCP tools through the `langroid.agent.tools.mcp` module. Here are the key patterns and approaches:

### MCP Tool Creation Methods

#### 1. Using the `@mcp_tool` Decorator (Module Level)
```python
from langroid.agent.tools.mcp import mcp_tool
from fastmcp.client.transports import StdioTransport

transport = StdioTransport(command="...", args=[...])

@mcp_tool(transport, "tool_name")
class MyTool(lr.ToolMessage):
    async def handle_async(self):
        result = await self.call_tool_async()
        # custom processing
        return result
```

**Important**: The decorator creates the transport connection at module import time, so it must be used at module level (not inside async functions).

#### 2. Using `get_tool_async` (Inside Async Functions)
```python
from langroid.agent.tools.mcp.fastmcp_client import get_tool_async

async def main():
    transport = StdioTransport(command="...", args=[...])
    BaseTool = await get_tool_async(transport, "tool_name")
    
    class MyTool(BaseTool):
        async def handle_async(self):
            result = await self.call_tool_async()
            # custom processing
            return result
```

**Use this approach when**:
- Creating tools inside async functions
- Need to avoid event loop conflicts
- Want to delay transport creation until runtime

### Transport Types and Event Loop Considerations

- **StdioTransport**: Creates subprocess immediately, can cause "event loop closed" errors if created at module level in certain contexts
- **SSETransport**: HTTP-based, generally safer for module-level creation
- **Best Practice**: Create transports inside async functions when possible, use `asyncio.run()` wrapper for Fire CLI integration

### Tool Message Request Field and Agent Handlers

When you get an MCP tool named "my_tool", Langroid automatically:

1. **Sets the `request` field**: The dynamically created ToolMessage subclass has `request = "my_tool"`
2. **Enables custom agent handlers**: Agents can define these methods:
   - `my_tool()` - synchronous handler
   - `my_tool_async()` - async handler

The agent's message routing system automatically calls these handlers when the tool is used.

### Custom `handle_async` Method Override

Both decorator and non-decorator approaches support overriding `handle_async`:

```python
class MyTool(BaseTool):  # or use @mcp_tool decorator
    async def handle_async(self):
        # Get raw result from MCP server
        result = await self.call_tool_async()
        
        # Option 1: Return processed result to LLM (continues conversation)
        return f"<ProcessedResult>{result}</ProcessedResult>"
        
        # Option 2: Return ResultTool to terminate task
        return MyResultTool(answer=result)
```

### Common Async Issues and Solutions

**Problem**: "RuntimeError: asyncio.run() cannot be called from a running event loop"
**Solution**: Use `get_tool_async` instead of `@mcp_tool` decorator when already in async context

**Problem**: "RuntimeError: Event loop is closed"
**Solution**: 
- Move transport creation inside async functions
- Use `asyncio.run()` wrapper for Fire CLI integration:
```python
if __name__ == "__main__":
    import asyncio
    def run_main(**kwargs):
        asyncio.run(main(**kwargs))
    Fire(run_main)
```

### MCP Tool Integration Examples

See `examples/mcp/` for working examples:
- `gitmcp.py` - HTTP-based SSE transport
- `pyodide_code_executor.py` - Subprocess-based stdio transport with proper async handling

## Testing and Tool Message Patterns

### MockLM for Testing Tool Generation
- Use `MockLM` with `response_dict` to simulate LLM responses that include tool messages
- Set `tools=[ToolClass]` or `enable_message=[ToolClass]` on the agent to enable tool handling
- The `try_get_tool_messages()` method can extract tool messages from LLM responses with `all_tools=True`

### Task Termination Control
- `TaskConfig` has `done_if_tool` parameter to terminate tasks when any tool is generated
- `Task.done()` method checks `result.agent_response` for tool content when this flag is set
- Useful for workflows where tool generation signals task completion

### Testing Tool-Based Task Flows
```python
# Example: Test task termination on tool generation
config = TaskConfig(done_if_tool=True)
task = Task(agent, config=config)
response_dict = {"content": '{"request": "my_tool", "param": "value"}'}
```

## Multi-Agent System Development

### Important Patterns and Best Practices

#### 1. Pydantic Imports
**ALWAYS import Pydantic classes from `langroid.pydantic_v1`**, not from `pydantic` directly:
```python
# CORRECT
from langroid.pydantic_v1 import Field, BaseModel

# WRONG - will cause issues
from pydantic import Field, BaseModel
```

#### 2. Tool Name References in System Messages
When referencing tool names in f-strings within system messages, use the `.name()` method:
```python
system_message: str = f"""
Use {MyTool.name()} to perform the action.
"""
```
This works at module level in configs, but be aware that complex initialization at module level can sometimes cause issues.

#### 3. Agent Configuration with LLM
Always specify the LLM configuration explicitly in agent configs:
```python
class MyAgentConfig(lr.ChatAgentConfig):
    name: str = "MyAgent"
    llm: lm.OpenAIGPTConfig = lm.OpenAIGPTConfig(
        chat_model="gpt-4",  # or "gpt-4.1" etc.
    )
    system_message: str = "..."
```

#### 4. Tool Organization in Multi-Agent Systems
When tools delegate to agents:
- Define agent configs and agents BEFORE the tools that use them
- Tools can directly instantiate agents in their `handle()` methods:
```python
class MyTool(lr.ToolMessage):
    def handle(self) -> str:
        agent = MyAgent(MyAgentConfig())
        task = lr.Task(agent, interactive=False)
        result = task.run(prompt)
        return result.content
```

#### 5. Task Termination with Done Sequences
Use `done_sequences` for precise task termination control:
```python
# For a task that should complete after: Tool -> Agent handles -> LLM responds
task = lr.Task(
    agent,
    interactive=False,
    config=lr.TaskConfig(done_sequences=["T,A,L"]),
)
```

Common patterns:
- `"T,A"` - Tool used and handled by agent
- `"T,A,L"` - Tool used, handled, then LLM responds
- `"T[specific_tool],A"` - Specific tool used and handled

See `docs/notes/task-termination.md` for comprehensive documentation.

#### 6. Handling Non-Tool LLM Responses
Use `handle_llm_no_tool` in agent configs to handle cases where the LLM forgets to use a tool:
```python
class MyAgentConfig(lr.ChatAgentConfig):
    handle_llm_no_tool: str = "You FORGOT to use one of your TOOLs!"
```

#### 7. Agent Method Parameters
Note that `ChatAgentConfig` does not have a `use_tools` parameter. Instead, enable tools on the agent after creation:
```python
agent = MyAgent(config)
agent.enable_message([Tool1, Tool2, Tool3])  # Pass list of tool classes
```

## Commit and Pull Request Guidelines

- Never include "co-authored by Claude Code" or "created by Claude" in commit messages or pull request descriptions

## Codecov Badge Fix (June 2025)

- Fixed broken Codecov badge in README by removing the token parameter from the URL
- Changed from `https://codecov.io/gh/langroid/langroid/branch/main/graph/badge.svg?token=H94BX5F0TE` to `https://codecov.io/gh/langroid/langroid/graph/badge.svg`
- Tokens are not needed for public repositories and can cause GitHub rendering issues
</file>

<file path="ai-notes/repomix-plan.md">
## Plan to Add llms-no-tests.txt

### Overview
Create a third version of the repomix output that excludes all test files from the `tests/` directory. This will provide a more concise version focused only on source code without test implementations.

### Steps:

1. **Create ai-scratchpads directory and save this plan** ✓
   - Create directory: `mkdir -p ai-scratchpads`
   - Save this plan to `ai-scratchpads/repomix-plan.md`

2. **Create temporary repomix configuration**
   - Copy existing `repomix.config.json` to `repomix-no-tests.config.json`
   - Add `"tests/**"` to the `customPatterns` array in the `ignore` section
   - Add `"llms-no-tests.txt"` to the ignore patterns to prevent recursive inclusion

3. **Generate the new output file**
   - Run: `repomix --config repomix-no-tests.config.json -o llms-no-tests.txt`
   - This will create a new file excluding all test files

4. **Clean up and update documentation**
   - Remove the temporary `repomix-no-tests.config.json` file
   - Update `ai-instructions/claude-repomix-instructions.md` to mention the third variant
   - Add a note about generating the no-tests version with the command:
     ```bash
     # No-tests version (excludes tests directory)
     repomix --config repomix-no-tests.config.json -o llms-no-tests.txt
     ```

### Expected Result
- A new file `llms-no-tests.txt` that contains all source code except test files
- This will be smaller than the standard `llms.txt` but larger than `llms-compressed.txt`
- Useful for LLM analysis when test implementations are not needed

### File Size Expectations
Based on the current setup:
- `llms.txt`: ~3.3 MB (782K tokens)
- `llms-compressed.txt`: ~1.6 MB (434K tokens)
- `llms-no-tests.txt`: Expected to be between these sizes, excluding test code

## Results and Conclusions

### Actual Token Counts
After generating all variants, here are the actual token counts:
- `llms.txt`: 782K tokens (standard version with tests)
- `llms-compressed.txt`: 434K tokens (compressed version with tests)
- `llms-no-tests.txt`: 652K tokens (no tests version)
- `llms-no-tests-compressed.txt`: 400K tokens (compressed no-tests version)

### Key Observations
1. **Limited Impact of Excluding Tests**: Removing test files only reduced tokens by ~130K (17% reduction), suggesting that test files don't constitute a major portion of the codebase.

2. **Compression More Effective**: The compression feature provides a much more significant reduction (~45-50% reduction) compared to just excluding tests.

3. **Minimal Benefit of Combined Approach**: The compressed no-tests version (400K) is only marginally smaller than the compressed version with tests (434K) - a difference of just 34K tokens or ~8%.

### Recommendations
- For most use cases, the standard `llms-compressed.txt` (434K tokens) is likely sufficient
- The no-tests variants might be useful for specific scenarios where test implementation details would confuse the LLM or are explicitly not needed
- The marginal benefit of excluding tests doesn't justify maintaining multiple variants unless there's a specific need

### Files Created
- `repomix-no-tests.config.json` - Permanent config file for generating no-tests versions
- `llms-no-tests.txt` - Full version without tests (652K tokens)
- `llms-no-tests-compressed.txt` - Compressed version without tests (400K tokens)
</file>

<file path="docs/blog/posts/chat-completion.md">
---
title: 'Language Models: Completion and Chat-Completion'
draft: false
date: 2023-09-19
authors: 
  - pchalasani
categories:
  - langroid
  - llm
  - local-llm
  - chat
comments: true
---

Transformer-based language models are fundamentally next-token predictors, so 
naturally all LLM APIs today at least provide a completion endpoint. 
If an LLM is a next-token predictor, how could it possibly be used to 
generate a response to a question or instruction, or to engage in a conversation with 
a human user? This is where the idea of "chat-completion" comes in.
This post is a refresher on the distinction between completion and chat-completion,
and some interesting details on how chat-completion is implemented in practice.

<!-- more -->

## Language Models as Next-token Predictors

A Language Model is essentially a "next-token prediction" model,
and so all LLMs today provide a "completion" endpoint, typically something like:
`/completions` under the base URL.

The endpoint simply takes a prompt and returns a completion (i.e. a continuation).

A typical prompt sent to a completion endpoint might look like this:
```
The capital of Belgium is 
```
and the LLM will return a completion like this:
```
Brussels.
```
OpenAI's GPT3 is an example of a pure completion LLM.
But interacting with a completion LLM is not very natural or useful:
you cannot give instructions or ask questions; instead you would always need to 
formulate your input as a prompt whose natural continuation is your desired output.
For example, if you wanted the LLM to highlight all proper nouns in a sentence,
you would format it as the following prompt:

**Chat-To-Prompt Example:** Chat/Instruction converted to a completion prompt.

```
User: here is a sentence, the Assistant's task is to identify all proper nouns.
     Jack lives in Bosnia, and Jill lives in Belgium.
Assistant:    
```
The natural continuation of this prompt would be a response listing the proper nouns,
something like:
```
John, Bosnia, Jill, Belgium are all proper nouns.
```

This _seems_ sensible in theory, but a "base" LLM that performs well on completions
may _not_ perform well on these kinds of prompts. The reason is that during its training, it may not
have been exposed to very many examples of this type of prompt-response pair.
So how can an LLM be improved to perform well on these kinds of prompts?

## Instruction-tuned, Aligned LLMs 

This brings us to the heart of the innovation behind the wildly popular ChatGPT:
it uses an enhancement of GPT3 that (besides having a lot more parameters),
was _explicitly_ fine-tuned on instructions (and dialogs more generally) -- this is referred to
as **instruction-fine-tuning** or IFT for short. In addition to fine-tuning instructions/dialogs,
the models behind ChatGPT (i.e., GPT-3.5-Turbo and GPT-4) are further tuned to produce
responses that _align_ with human preferences (i.e. produce responses that are more helpful and safe),
using a procedure called Reinforcement Learning with Human Feedback (RLHF).
See this [OpenAI InstructGPT Paper](https://arxiv.org/pdf/2203.02155.pdf) for details on these techniques and references to the 
original papers that introduced these ideas. Another recommended read is Sebastian 
Raschka's post on [RLHF and related techniques](https://magazine.sebastianraschka.com/p/llm-training-rlhf-and-its-alternatives). 

For convenience, we refer to the combination of IFT and RLHF as **chat-tuning**.
A chat-tuned LLM can be expected to perform well on prompts such as the one in 
the Chat-To-Prompt Example above. These types of prompts are still unnatural, however, 
so as a convenience, chat-tuned LLM API servers also provide a "chat-completion" 
endpoint (typically `/chat/completions` under the base URL), which allows the user
to interact with them in a natural dialog, which might look like this
(the portions in square brackets are indicators of who is generating the text):

```
[User] What is the capital of Belgium?
[Assistant] The capital of Belgium is Brussels.
```
or
```
[User] In the text below, find all proper nouns:
    Jack lives in Bosnia, and Jill lives in Belgium.
[Assistant] John, Bosnia, Jill, Belgium are all proper nouns.
[User] Where does John live?
[Assistant] John lives in Bosnia.
```

## Chat Completion Endpoints: under the hood

How could this work, given that LLMs are fundamentally next-token predictors?
This is a convenience provided by the LLM API service (e.g. from OpenAI or
local model server libraries):
when a user invokes the chat-completion endpoint (typically
at `/chat/completions` under the base URL), under the hood, the server converts the
instructions and multi-turn chat history into a single string, with annotations indicating
user and assistant turns, and ending with something like "Assistant:"
as in the Chat-To-Prompt Example above.

Now the subtle detail to note here is this:

>It matters _how_ the
dialog (instructions plus chat history) is converted into a single prompt string.
Converting to a single prompt by simply concatenating the
instructions and chat history using an "intuitive" format (e.g. indicating
user, assistant turns using "User", "Assistant:", etc.) _can_ work,
however most local LLMs are trained on a _specific_ prompt format.
So if we format chats in a different way, we may get odd/inferior results.

## Converting Chats to Prompts: Formatting Rules

For example, the llama2 models are trained on a format where the user's input is bracketed within special strings `[INST]`
and `[/INST]`. There are other requirements that we don't go into here, but
interested readers can refer to these links:

- A reddit thread on the [llama2 formats](https://www.reddit.com/r/LocalLLaMA/comments/155po2p/get_llama_2_prompt_format_right/)
- Facebook's [llama2 code](https://github.com/facebookresearch/llama/blob/main/llama/generation.py#L44)
- Langroid's [llama2 formatting code](https://github.com/langroid/langroid/blob/main/langroid/language_models/prompt_formatter/llama2_formatter.py)

A dialog fed to a Llama2 model in its expected prompt format would look like this:

```
<s>[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

Hi there! 
[/INST] 
Hello! How can I help you today? </s>
<s>[INST] In the text below, find all proper nouns:
    Jack lives in Bosnia, and Jill lives in Belgium.
 [/INST] 
John, Bosnia, Jill, Belgium are all proper nouns. </s><s> 
[INST] Where does Jack live? [/INST] 
Jack lives in Bosnia. </s><s>
[INST] And Jill? [/INST]
Jill lives in Belgium. </s><s>
[INST] Which are its neighboring countries? [/INST]
```

This means that if an LLM server library wants to provide a chat-completion endpoint for
a local model, it needs to provide a way to convert chat history to a single prompt
using the specific formatting rules of the model.
For example the [`oobabooga/text-generation-webui`](https://github.com/oobabooga/text-generation-webui) 
library has an extensive set of chat formatting [templates](https://github.com/oobabooga/text-generation-webui/tree/main/instruction-templates)
for a variety of models, and their model server auto-detects the
format template from the model name.

!!! note "Chat completion model names: look for 'chat' or 'instruct' in the name"
    You can search for a variety of models on the [HuggingFace model hub](https://huggingface.co/models).
    For example if you see a name `Llama-2-70B-chat-GGUF` you know it is chat-tuned.
    Another example of a chat-tuned model is `Llama-2-7B-32K-Instruct` 
    
A user of these local LLM server libraries thus has two options when using a 
local model in chat mode:

- use the _chat-completion_ endpoint, and let the underlying library handle the chat-to-prompt formatting, or
- first format the chat history according to the model's requirements, and then use the
  _completion_ endpoint

## Using Local Models in Langroid

Local models can be used in Langroid by defining a `LocalModelConfig` object.
More details are in this [tutorial](https://langroid.github.io/langroid/blog/2023/09/14/using-langroid-with-local-llms/), 
but here we briefly discuss prompt-formatting in this context.
Langroid provides a built-in [formatter for LLama2 models](https://github.com/langroid/langroid/blob/main/langroid/language_models/prompt_formatter/llama2_formatter.py), 
so users looking to use llama2 models with langroid can try either of these options, by setting the
`use_completion_for_chat` flag in the `LocalModelConfig` object
(See the local-LLM [tutorial](https://langroid.github.io/langroid/blog/2023/09/14/using-langroid-with-local-llms/) for details).

When this flag is set to `True`, the chat history is formatted using the built-in 
Langroid llama2 formatter and the completion endpoint are used. When the flag is set to `False`, the chat 
history is sent directly to the chat-completion endpoint, which internally converts the 
chat history to a prompt in the expected llama2 format.

For local models other than Llama2, users can either:

- write their own formatters by writing a class similar to `Llama2Formatter` and 
then setting the `use_completion_for_chat` flag to `True` in the `LocalModelConfig` object, or
- use an LLM server library (such as the `oobabooga` library mentioned above) that provides a chat-completion endpoint, 
_and converts chats to single prompts under the hood,_ and set the
  `use_completion_for_chat` flag to `False` in the `LocalModelConfig` object.

You can use a similar approach if you are using an LLM application framework other than Langroid.


<iframe src="https://langroid.substack.com/embed" width="480" height="320" style="border:1px solid #EEE; background:white;" frameborder="0" scrolling="no"></iframe>
</file>

<file path="docs/blog/posts/langroid-architecture.md">
---
title: "Overview of Langroid's Multi-Agent Architecture (prelim)"
draft: false
date: 2024-08-15
authors:
- pchalasani
- nils
- jihye
- someshjha
categories:
- langroid
- multi-agent
- llm
comments: true
---


## Agent, as an intelligent message transformer

A natural and convenient abstraction in designing a complex
LLM-powered system is the notion of an *agent* that is instructed to be responsible for a specific aspect of the 
overall task. In terms of code, an *Agent* is essentially a class representing an intelligent entity that can 
respond to *messages*, i.e., an agent is simply a *message transformer*.
An agent typically encapsulates an (interface to an) LLM, and may also be equipped with so-called *tools* (as 
described below) and *external documents/data* (e.g., via a vector database, as described below).
Much like a team of humans, agents interact by exchanging messages, in a manner reminiscent of the 
[*actor framework*](https://en.wikipedia.org/wiki/Actor_model) in programming languages.
An *orchestration mechanism* is needed to manage the flow of messages between agents, to ensure that progress is 
made towards completion of the task, and to handle the inevitable cases where an agent deviates from instructions.
Langroid is founded on this *multi-agent programming* paradigm, where agents are 
first-class citizens, acting as message transformers, and communicate by exchanging messages.

<!-- more -->

To build useful applications with LLMs, we need to endow them with the ability to
trigger actions (such as API calls, computations, database queries, etc) or send structured messages to other agents 
or downstream processes. *Tools* provide these capabilities, described next.

## Tools, also known as functions

An LLM is essentially a text transformer; i.e.,  in response to some input text, 
it produces a text response. Free-form text responses are ideal when we want to generate a description, answer, or summary for human consumption, or even a question for another agent to answer.
However, in some cases, we would like the responses to be more structured, for example 
to trigger external *actions* (such as an API call, code execution, or a database query),
or for unambiguous/deterministic handling by a downstream process or another agent. 
In such cases, we would instruct the LLM to produce a *structured* output, typically in JSON format, with various 
pre-specified fields, such as code, an SQL query, parameters of an API call, and so on. These structured responses 
have come to be known as *tools*, and the LLM is said to *use* a tool when it produces a structured response 
corresponding to a specific tool. To elicit a tool response from an LLM, it needs to be instructed on the expected tool format and the conditions under which it should use the tool.
To actually use a tool emitted by an LLM, a *tool handler* method must be defined as well.
The tool handler for a given tool is triggered when it is recognized in the LLM's response.

### Tool Use: Example

As a simple example, a SQL query tool can be specified as a JSON structure with a `sql` 
field (containing the SQL query) and a `db` field (containing the name of the database).
The LLM may be instructed with a system prompt of the form:
> When the user asks a question about employees, use the SQLTool described in the below schema,
> and the results of this tool will be sent back to you, and you can use these to respond to
> the user's question, or correct your SQL query if there is a syntax error.

The tool handler would detect this specific tool in the LLM's response, parse this JSON structure, 
extract the `sql` and `db` fields, run the query on the specified database, 
and return the result if the query ran successfully, otherwise return an error message.
Depending on how the multi-agent system is organized, the query result or error message may be handled by the same agent
(i.e., its LLM), which may either summarize the results in narrative form, or revise the query if the error message 
indicates a syntax error.

## Agent-oriented programming: Function-Signatures

If we view an LLM as a function with signature `string -> string`,
it is possible to express the concept of an agent, tool, and other constructs
in terms of derived function signatures, as shown in the table below.
Adding `tool` (or function calling) capability to an LLM requires a parser (that recognizes 
that the LLM has generated a tool) and a callback that performs arbitrary computation and returns a string.
The serialized instances of tools `T` correspond to a language `L`; 
Since by assumption, the LLM is capable of producing outputs in $L$, 
this allows the LLM to express the intention to execute a Callback with arbitrary instances 
of `T`. In the last row, we show how an Agent can be viewed as a function signature
involving its state `S`.


| Function Description | Function Signature                                                                                                |
|----------------------|-------------------------------------------------------------------------------------------------------------------|
| LLM | `[Input Query] -> string` <br> `[Input Query]` is the original query.                                             |
| Chat interface | `[Message History] x [Input Query] -> string` <br> `[Message History]` consists of  previous messages[^1].        |
| Agent | `[System Message] x [Message History] x [Input Query] -> string` <br> `[System Message]` is the system prompt. |
| Agent with tool | `[System Message] x (string -> T) x (T -> string) x [Message History] x [Input Query] -> string`                  |
| Parser with type `T` | `string -> T`                                                                                                     |
| Callback with type `T` | `T -> string`                                                                                                     |
| General Agent with state type `S` | `S x [System Message] x (string -> T) x (S x T -> S x string) x [Message History] x [Input Query] -> S x string`  |

[^1]: Note that in reality, separator tokens are added to distinguish messages, and the messages are tagged with metadata indicating the sender, among other things.

## Multi-Agent Orchestration


### An Agent's "Native" Responders

When building an LLM-based multi-agent system, an orchestration mechanism is critical to manage the flow of messages 
between agents, to ensure task progress, and handle inevitable LLM deviations from instructions. Langroid provides a 
simple yet versatile orchestration mechanism that seamlessly handles:

- user interaction,
- tool handling,
- sub-task delegation

We view an agent as a message transformer; 
it may transform an incoming message using one of its three "native" responder methods, all of which have the same 
function signature: `string -> string`. These methods are:

- `llm_response` returns the LLM's response to the input message.
Whenever this method is invoked, the agent updates its dialog history (typically consisting of alternating user and LLM messages).
- `user_response` prompts the user for input and returns their response.
- `agent_response` by default only handles a `tool message` (i.e., one that contains an llm-generated structured 
response): it performs any requested actions, and returns the result as a string. An `agent_response` method can have 
other uses besides handling tool messages, such as handling scenarios where an LLM ``forgot'' to use a tool, 
or used a tool incorrectly, and so on.

To see why it is useful to have these responder methods, consider first a simple example of creating a basic chat loop
with the user. It is trivial to create such a loop by alternating between `user_response` and `llm_response`. 
Now suppose we instruct the agent to either directly answer the user's question or perform a web-search. Then it is possible that
sometimes the `llm_response` will produce a "tool message", say `WebSearchTool`, which we would handle with the
`agent_response` method. This requires a slightly different, and more involved, way of iterating among the agent's
responder methods. 

### Tasks: Encapsulating Agent Orchestration

From a coding perspective, it is useful to hide the actual iteration logic by wrapping an Agent class
in a separate class, which we call a `Task`, which encapsulates all of the orchestration logic. Users of the Task class
can then define the agent, tools, and any sub-tasks, wrap the agent in a task object of class Task, and simply call
`task.run()`, letting the Task class deal with the details of orchestrating the agent's responder methods, determining
task completion, and invoking sub-tasks.

### Responders in a Task: Agent's native responders and sub-tasks

The orchestration mechanism of a `Task` object works as follows. When a `Task` object is created from an agent, a 
sequence of eligible responders is created, which includes the agent's three "native" responder agents in the sequence:
`agent_response`, `llm_response`, `user_response`. 
The type signature of the task's run method is `string -> string`, just like the Agent's
native responder methods, and this is the key to seamless delegation of tasks to sub-tasks. A list of subtasks can be
added to a `Task` object via `task.add_sub_tasks([t1, t2, ... ])`, where `[t1, t2, ...]` are other 
`Task` objects. The result of this is that the run method of each sub-task is appended to the sequence of eligible 
responders in the parent task object.

### Task Orchestration: Updating the Current Pending Message (CPM)

A task always maintains a *current pending message* (CPM), which is the latest message "awaiting" a valid response 
from a responder, which updates the CPM. 
At a high level the `run` method of a task attempts to repeatedly find a valid response to the 
CPM, until the task is done. (Note that this paradigm is somewhat reminescent of a *Blackboard* architecture, where
agents take turns deciding whether they can update the shared message on the "blackboard".)
This is achieved by repeatedly invoking the `step` method, which represents a "turn" in the conversation.
The `step` method sequentially tries the eligible responders from the beginning of the eligible-responders list, until it
finds a valid response, defined as a non-null or terminating message (i.e. one that signals that the task is done). In
particular, this `step()` algorithm implies that a Task delegates (or "fails over") to a sub-task only if the task's 
native responders have no valid response. 

There are a few simple rules that govern how `step` works: 

- a responder entity (either a sub-task or a native entity -- one of LLM, Agent, or User) cannot 
  respond if it just responded in the previous step (this prevents a responder from "talking to itself". 
- when a response signals that the task is done (via a `DoneTool` or a "DONE" string) the task is ready to exit and 
  return the CPM as the result of the task. 
- when an entity "in charge" of the task has a null response, the task is considered finished and ready to exit.
- if the response of an entity or subtask is a structured message containing a recipient field, then the specified recipient task or entity will
be the only one eligible to respond at the next step.

Once a valid response is found in a step, the CPM is updated to this response, and the next step starts the search for a
valid response from the beginning of the eligible responders list. When a response signals that the task is done, 
the run method returns the CPM as the result of the task. This is a highly
simplified account of the orchestration mechanism, and the actual implementation is more involved.

The above simple design is surprising powerful and can support a wide variety of task structures, including trees and
DAGs. As a simple illustrative example, tool-handling has a natural implementation. The LLM is instructed to use a
certain JSON-structured message as a tool, and thus the `llm_response` method can produce a structured message, such 
as an SQL query.  This structured message is then handled by the `agent_response` method, and the resulting message updates the CPM. The
`llm_response` method then becomes eligible to respond again: for example if the agent's response contains an SQL 
error, the LLM would retry its query, and if the agent's response consists of the query results, the LLM would
respond with a summary of the results.

The Figure below depicts the task orchestration and delegation mechanism,
showing how iteration among responder methods works when a  Task `T` has sub-tasks `[T1, T2]` and `T1` has a 
sub-task `T3`. 


![langroid-arch.png](figures/langroid-arch.png)
</file>

<file path="docs/blog/posts/langroid-intro.md">
---
title: 'Langroid: Harness LLMs with Multi-Agent Programming'
draft: false
date: 2023-09-03
authors: 
  - pchalasani
categories:
  - langroid
  - llm
comments: true
---

# Langroid: Harness LLMs with Multi-Agent Programming

## The LLM Opportunity

Given the remarkable abilities of recent Large Language Models (LLMs), there
is an unprecedented opportunity to build intelligent applications powered by
this transformative technology. The top question for any enterprise is: how
best to harness the power of LLMs for complex applications? For technical and
practical reasons, building LLM-powered applications is not as simple as
throwing a task at an LLM-system and expecting it to do it.

<!-- more -->


## Langroid's Multi-Agent Programming Framework

Effectively leveraging LLMs at scale requires a *principled programming
framework*. In particular, there is often a need to maintain multiple LLM
conversations, each instructed in different ways, and "responsible" for
different aspects of a task.


An *agent* is a convenient abstraction that encapsulates LLM conversation
state, along with access to long-term memory (vector-stores) and tools (a.k.a functions
or plugins). Thus a **Multi-Agent Programming** framework is a natural fit
for complex LLM-based applications.

> Langroid is the first Python LLM-application framework that was explicitly
designed  with Agents as first-class citizens, and Multi-Agent Programming
as the core  design principle. The framework is inspired by ideas from the
[Actor Framework](https://en.wikipedia.org/wiki/Actor_model).

Langroid allows an intuitive definition of agents, tasks and task-delegation
among agents. There is a principled mechanism to orchestrate multi-agent
collaboration. Agents act as message-transformers, and take turns responding to (and
transforming) the current message. The architecture is lightweight, transparent,
flexible, and allows other types of orchestration to be implemented.
Besides Agents, Langroid also provides simple ways to directly interact with LLMs and vector-stores.


## Highlights
- **Agents as first-class citizens:** The `Agent` class encapsulates LLM conversation state,
  and optionally a vector-store and tools. Agents are a core abstraction in Langroid;
  Agents act as _message transformers_, and by default provide 3 _responder_ methods, one corresponding to each
  entity: LLM, Agent, User.
- **Tasks:** A Task class wraps an Agent, gives the agent instructions (or roles, or goals),
  manages iteration over an Agent's responder methods,
  and orchestrates multi-agent interactions via hierarchical, recursive
  task-delegation. The `Task.run()` method has the same
  type-signature as an Agent's responder's methods, and this is key to how
  a task of an agent can delegate to other sub-tasks: from the point of view of a Task,
  sub-tasks are simply additional responders, to be used in a round-robin fashion
  after the agent's own responders.
- **Modularity, Reusability, Loose coupling:** The `Agent` and `Task` abstractions allow users to design
  Agents with specific skills, wrap them in Tasks, and combine tasks in a flexible way.
- **LLM Support**: Langroid supports OpenAI LLMs including GPT-3.5-Turbo,
  GPT-4.
- **Caching of LLM prompts, responses:** Langroid by default uses [Redis](https://redis.com/try-free/) for caching.
- **Vector-stores**: [Qdrant](https://qdrant.tech/), [Chroma](https://www.trychroma.com/), LanceDB, Pinecone, PostgresDB (PGVector), Weaviate are currently supported.
  Vector stores allow for Retrieval-Augmented-Generaation (RAG).
- **Grounding and source-citation:** Access to external documents via vector-stores
  allows for grounding and source-citation.
- **Observability, Logging, Lineage:** Langroid generates detailed logs of multi-agent interactions and
  maintains provenance/lineage of messages, so that you can trace back
  the origin of a message.
- **Tools/Plugins/Function-calling**: Langroid supports OpenAI's recently
  released [function calling](https://platform.openai.com/docs/guides/gpt/function-calling)
  feature. In addition, Langroid has its own native equivalent, which we
  call **tools** (also known as "plugins" in other contexts). Function
  calling and tools have the same developer-facing interface, implemented
  using [Pydantic](https://docs.pydantic.dev/latest/),
  which makes it very easy to define tools/functions and enable agents
  to use them. Benefits of using Pydantic are that you never have to write
  complex JSON specs for function calling, and when the LLM
  hallucinates malformed JSON, the Pydantic error message is sent back to
  the LLM so it can fix it!

<iframe src="https://langroid.substack.com/embed" width="480" height="320" style="border:1px solid #EEE; background:white;" frameborder="0" scrolling="no"></iframe>
</file>

<file path="docs/blog/posts/langroid-knowledge-graph.md">
---
title: 'Langroid: Knolwedge Graph RAG powered by Neo4j'
draft: false
date: 2024-01-18
authors: 
  - mohannad
categories:
  - langroid
  - neo4j
  - rag
  - knowledge-graph
comments: true
---

## "Chat" with various sources of information
LLMs are increasingly being used to let users converse in natural language with 
a variety of types of data sources:
<!-- more -->
- unstructured text documents: a user's query is augmented with "relevant" documents or chunks
  (retrieved from an embedding-vector store) and fed to the LLM to generate a response -- 
  this is the idea behind Retrieval Augmented Generation (RAG).
- SQL Databases: An LLM translates a user's natural language question into an SQL query,
  which is then executed by another module, sending results to the LLM, so it can generate
  a natural language response based on the results.
- Tabular datasets: similar to the SQL case, except instead of an SQL Query, the LLM generates 
  a Pandas dataframe expression.

Langroid has had specialized Agents for the above scenarios: `DocChatAgent` for RAG with unstructured
text documents, `SQLChatAgent` for SQL databases, and `TableChatAgent` for tabular datasets.

## Adding support for Neo4j Knowledge Graphs

Analogous to the SQLChatAgent, Langroid now has a 
[`Neo4jChatAgent`](https://github.com/langroid/langroid/blob/main/langroid/agent/special/neo4j/neo4j_chat_agent.py) 
to interact with a Neo4j knowledge graph using natural language.
This Agent has access to two key tools that enable it to handle a user's queries:

- `GraphSchemaTool` to get the schema of a Neo4j knowledge graph.
- `CypherRetrievalTool` to generate Cypher queries from a user's query.
Cypher is a specialized query language for Neo4j, and even though it is not as widely known as SQL,
most LLMs today can generate Cypher Queries.

Setting up a basic Neo4j-based RAG chatbot is straightforward. First ensure 
you set these environment variables (or provide them in a `.env` file):
```bash
NEO4J_URI=<uri>
NEO4J_USERNAME=<username>
NEO4J_PASSWORD=<password>
NEO4J_DATABASE=<database>
```

Then you can configure and define a `Neo4jChatAgent` like this:
```python
import langroid as lr
import langroid.language_models as lm

from langroid.agent.special.neo4j.neo4j_chat_agent import (
    Neo4jChatAgent,
    Neo4jChatAgentConfig,
    Neo4jSettings,
)

llm_config = lm.OpenAIGPTConfig()

load_dotenv()

neo4j_settings = Neo4jSettings()

kg_rag_agent_config = Neo4jChatAgentConfig(
    neo4j_settings=neo4j_settings,
    llm=llm_config, 
)
kg_rag_agent = Neo4jChatAgent(kg_rag_agent_config)
kg_rag_task = lr.Task(kg_rag_agent, name="kg_RAG")
kg_rag_task.run()
```


## Example: PyPi Package Dependency Chatbot

In the Langroid-examples repository, there is an example python 
[script](https://github.com/langroid/langroid-examples/blob/main/examples/kg-chat/)
showcasing tools/Function-calling + RAG using a `DependencyGraphAgent` derived from [`Neo4jChatAgent`](https://github.com/langroid/langroid/blob/main/langroid/agent/special/neo4j/neo4j_chat_agent.py).
This agent uses two tools, in addition to the tools available to `Neo4jChatAgent`:

- `GoogleSearchTool` to find package version and type information, as well as to answer 
 other web-based questions after acquiring the required information from the dependency graph.
- `DepGraphTool` to construct a Neo4j knowledge-graph modeling the dependency structure
   for a specific package, using the API at [DepsDev](https://deps.dev/).

In response to a user's query about dependencies, the Agent decides whether to use a Cypher query
or do a web search. Here is what it looks like in action:

<figure markdown>
  ![dependency-demo](../../assets/demos/dependency_chatbot.gif)
  <figcaption>
Chatting with the `DependencyGraphAgent` (derived from Langroid's `Neo4jChatAgent`).
When a user specifies a Python package name (in this case "chainlit"), the agent searches the web using
`GoogleSearchTool` to find the version of the package, and then uses the `DepGraphTool`
to construct the dependency graph as a neo4j knowledge graph. The agent then answers
questions by generating Cypher queries to the knowledge graph, or by searching the web.
  </figcaption>
</figure>
</file>

<file path="docs/blog/posts/langroid-lancedb.md">
---
title: 'Langroid: Multi-Agent Programming Framework for LLMs'
draft: true
date: 2024-01-10
authors: 
  - pchalasani
categories:
  - langroid
  - lancedb
  - rag
  - vector-database
comments: true
---

## Langroid: Multi-Agent Programming framework for LLMs

In this era of Large Language Models (LLMs), there is unprecedented demand to
create intelligent applications powered by this transformative technology. What
is the best way for developers to harness the potential of LLMs in complex
application scenarios? For a variety of technical and practical reasons (context
length limitations, LLM brittleness, latency, token-costs), this is not as
simple as throwing a task at an LLM system and expecting it to get done. What is
needed is a principled programming framework, offering the right set of
abstractions and primitives to make developers productive when building LLM
applications.
<!-- more -->
## Langroid's Elegant Multi-Agent Paradigm

The [Langroid](https://github.com/langroid/langroid) team (ex-CMU/UW-Madison researchers) 
has a unique take on this – they have built an open source Python framework to 
simplify LLM application development, using a Multi-Agent Programming paradigm. 
Langroid’s architecture is founded on Agents as first-class citizens: 
they are message-transformers, and accomplish tasks collaboratively via messages.

Langroid is emerging as a popular LLM framework; developers appreciate its clean
design and intuitive, extensible architecture. Programming with Langroid is
natural and even fun: you configure Agents and equip them with capabilities (
such as LLMs, vector-databases, Function-calling/tools), connect them and have
them collaborate via messages. This is a “Conversational Programming” paradigm,
and works with local/open and remote/proprietary LLMs. (Importantly, it does not
use LangChain or any other existing LLM framework).

<figure markdown>
  ![Langroid-card](../../assets/langroid-card-ossem-rust-1200x630.png){ width="800" }
  <figcaption>
An Agent serves as a convenient abstraction, encapsulating the state of LLM
conversations, access to vector stores, and various tools (functions or
plugins). A Multi-Agent Programming framework naturally aligns with the demands
of complex LLM-based applications.
</figcaption>
</figure>



## Connecting Agents via Tasks

In Langroid, a ChatAgent has a set of “responder” methods, one for each "entity":
an LLM, a human, and a tool-handler. However it does not have any way to iterate through
these responders. This is where the Task class comes in: A Task wraps an Agent
and gives it the ability to loop through its responders, via the `Task.run()` method. 

A Task loop is organized around simple rules that govern when a responder is eligible
to respond, what is considered a valid response, and when the task is complete.
The simplest example of a Task loop is an interactive chat with the human user. 
A Task also enables an Agent to interact with other agents: 
other tasks can be added to a task as sub-tasks, 
in a recursive, hierarchical (or DAG) structure. From a Task’s perspective,
sub-tasks are just additional responders, and present the same string-to-string 
message-transformation interface (function signature) as the Agent’s "native" responders. 
This is the key to composability of tasks in Langroid,
since a sub-task can act the same way as an Agent's "native" responders, and is subject
to the same rules of task orchestration. The result is that the same task orchestration
mechanism seamlessly enables tool handling, retries when LLM deviates, and 
delegation to sub-tasks. More details are in the Langroid [quick-start guide](https://langroid.github.io/langroid/quick-start/)

## A Taste of Coding with Langroid

To get started with Langroid, simply install it from pypi into your virtual environment:

```bash
pip install langroid
```
To directly chat with an OpenAI LLM, define the LLM configuration,
instantiate a language model object and interact with it:
(Langroid works with non-OpenAI local/propreitary LLMs as well,
see their [tutorial](https://langroid.github.io/langroid/tutorials/non-openai-llms/)) 
For the examples below, ensure you have a file `.env` containing your OpenAI API key
with this line: `OPENAI_API_KEY=sk-...`.
    
```python
import langroid as lr
import langroid.language_models as lm

llm_cfg = lm.OpenAIGPTConfig() # default GPT4-Turbo
mdl = lm.OpenAIGPT(llm_cfg)
mdl.chat("What is 3+4?", max_tokens=10)
```
The mdl does not maintain any conversation state; for that you need a `ChatAgent`:

```python
agent_cfg = lr.ChatAgentConfig(llm=llm_cfg)
agent = lr.ChatAgent(agent_cfg)
agent.llm_response("What is the capital of China?")
agent.llm_response("What about France?") # interprets based on previous msg
```
Wrap a ChatAgent in a Task to create a basic interactive loop with the user:

```python
task = lr.Task(agent, name="Bot")
task.run("Hello")
```
Have a Teacher Agent talk to a Student Agent:
    
```python
teacher = lr.ChatAgent(agent_cfg)
teacher_task = lr.Task(
    teacher, name="Teacher",
    system_message="""
        Ask your student simple number-based questions, and give feedback.
        Start with a question.
        """,
)
student = lr.ChatAgent(agent_cfg)
student_task = lr.Task(
    student, name="Student",
    system_message="Concisely answer your teacher's questions."
)
teacher_task.add_sub_task(student_task)
teacher_task.run()
```



## Retrieval Augmented Generation (RAG) and Vector Databases

One of the most popular LLM applications is question-answering 
on documents via Retrieval-Augmented Generation (RAG), powered by a vector database.
Langroid has a built-in DocChatAgent that incorporates a number of advanced RAG techniques, 
clearly laid out so they can be easily understood and extended.

### Built-in Support for LanceDB
<figure markdown>
  ![Langroid-lance](../../assets/langroid-lance.png){ width="800" }
  <figcaption>
Langroid uses LanceDB as the default vector store for its DocChatAgent.
</figcaption>
</figure>

Langroid's DocChatAgent uses the LanceDB serverless vector-database by default.
Since LanceDB uses file storage, it is easy to set up and use (no need for docker or cloud services),
and due to its use of the Lance columnar format, it is 
highly performant and scalable. 
In addition, Langroid has a specialized `LanceDocChatAgent` that leverages LanceDB's 
unique features such as Full-text search, SQL-like filtering, and pandas dataframe interop.
Setting up a basic RAG chatbot is as simple as (assume the previous imports):

```python
from langroid.agent.special.lance_doc_chat_agent import import (
    LanceDocChatAgent, DocChatAgentConfig
)
llm_config = lm.OpenAIGPTConfig()

rag_agent_config = DocChatAgentConfig(
    llm=llm_config, 
    doc_paths=["/path/to/my/docs"], # files, folders, or URLs.
)
rag_agent = LanceDocChatAgent(rag_agent_config)
rag_task = lr.Task(rag_agent, name="RAG")
rag_task.run()
```

For an example showcasing Tools/Function-calling + RAG in a multi-agent setup, see their quick-start
[Colab notebook](https://colab.research.google.com/github/langroid/langroid/blob/main/examples/Langroid_quick_start.ipynb)
which shows a 2-agent system where one agent is tasked with extracting structured information
from a document, and generates questions for the other agent to answer using RAG.
In the Langroid-examples repo there is a [script](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/chat_multi_extract.py) with the same functionality,
and here is what it looks like in action:


<figure markdown>
  ![lease-demo](../../assets/demos/lease-extractor-demo.gif){ width="800" }
  <figcaption>
Extracting structured info from a Commercial Lease using a 2-agent system, with 
a Tool/Function-calling and RAG. The Extractor Agent is told to extract information
in a certain structure, and it generates questions for the Document Agent
to answer using RAG.
</figcaption>
</figure>

## Retrieval Augmented Analytics

One of the unique features of LanceDB is its SQL-like filtering and Pandas dataframe interoperability.
LLMs are great at generating SQL queries, and also Pandas computation code such as `df.groupby("col").mean()`.
This opens up a very interesting possibility, which we call
**Retrieval Augmented Analytics:** Suppose a user has a large dataset of movie descriptions
with metadata such as rating, year and genre, and wants to ask:

> What is the highest-rated Comedy movie about college students made after 2010?

It is not hard to imagine that an LLM should be able to generate a **Query Plan** to answer this,
consisting of:

- A SQL-like filter: `genre = "Comedy" and year > 2010`
- A Pandas computation: `df.loc[df["rating"].idxmax()]`
- A rephrased query given the filter: "Movie about college students" (used for semantic/lexical search)

Langroid's Multi-Agent framework enables exactly this type of application. 
The [`LanceRAGTaskCreator`](https://github.com/langroid/langroid/blob/main/langroid/agent/special/lance_rag/lance_rag_task.py) takes a `LanceDocChatAgent` and adds two additional agents:

- QueryPlannerAgent: Generates the Query Plan
- QueryPlanCriticAgent: Critiques the Query Plan and Answer received from the RAG Agent, so that 
  the QueryPlanner can generate a better plan if needed.

Checkout the [`lance-rag-movies.py`](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/lance-rag-movies.py) script in the langroid-examples repo to try this out.

## Try it out and get involved!
This was just a glimpse of what you can do with Langroid and how your code would look.
Give it a shot and learn more about the features and roadmap of Langroid on their 
[GitHub repo](https://github.com/langroid/langroid). Langroid welcomes contributions,
and they have a friendly [Discord](https://discord.gg/ZU36McDgDs) community.

If you like it, don’t forget to drop a 🌟.
</file>

<file path="docs/blog/posts/local-llm-formatting.md">
---
title: 'Chat formatting in Local LLMs'
draft: true
date: 2024-01-25
authors: 
  - pchalasani
categories:
  - langroid
  - prompts
  - llm
  - local-llm
comments: true
---


In an (LLM performance) investigation, details matter!

And assumptions kill (your LLM performance).

I'm talking about chat/prompt formatting, especially when working with Local LLMs.

TL/DR -- details like chat formatting matter a LOT,
and trusting that the local LLM API is doing it correctly may be a mistake,
leading to inferior results.

<!-- more -->

🤔Curious? Here are some notes from the trenches when we built an app
(https://github.com/langroid/langroid/blob/main/examples/docqa/chat-multi-extract-local.py)
based entirely on a locally running Mistral-7b-instruct-v0.2  
(yes ONLY 7B parameters, compared to 175B+ for GPT4!)
that leverages Langroid Multi-agents, Tools/Function-calling and RAG to
reliably extract structured information from a document,
where an Agent is given a spec of the desired structure, and it generates
questions for another Agent to answer using RAG.

🔵LLM API types: generate and chat
LLMs are typically served behind two types of APIs endpoints:
⏺ A "generation" API, which accepts a dialog formatted as a SINGLE string, and
⏺ a "chat" API, which accepts the dialog as a LIST,
and as convenience formats it into a single string before sending to the LLM.

🔵Proprietary vs Local LLMs
When you use a proprietary LLM API (such as OpenAI or Claude), for convenience
you can use their "chat" API, and you can trust that it will format the dialog
history correctly (or else they wouldn't be in business!).

But with a local LLM, you have two choices of where to send the dialog history:
⏺ you could send it to the "chat" API and trust that the server will format it correctly,
⏺ or you could format it yourself and send it to the "generation" API.

🔵Example of prompt formatting?
Suppose your system prompt and dialog look like this:

System Prompt/Instructions: when I give you a number, respond with its double
User (You): 3
Assistant (LLM): 6
User (You): 9

Mistral-instruct models expect this chat to be formatted like this
(note that the system message is combined with the first user message):
"<s>[INST] when I give you a number, respond with its double 3 [/INST] 6 [INST] 9 [/INST]"

🔵Why does it matter?
It matters A LOT -- because each type of LLM (llama2, mistral, etc) has
been trained and/or fine-tuned on chats formatted in a SPECIFIC way, and if you
deviate from that, you may get odd/inferior results.

🔵Using Mistral-7b-instruct-v0.2 via oobabooga/text-generation-webui
"Ooba" is a great library (https://github.com/oobabooga/text-generation-webui)
that lets you spin up an OpenAI-like API server for
local models, such as llama2, mistral, etc. When we used its chat endpoint
for a Langroid Agent, we were getting really strange results,
with the LLM sometimes thinking it is the user! 😧

Digging in, we found that their internal formatting template was
wrong, and it was formatting the system prompt as if it's
the first user message -- this leads to the LLM interpreting the first user
message as an assistant response, and so on -- no wonder there was role confusion!

💥Langroid solution:
To avoid these issues, in Langroid we now have a formatter
(https://github.com/langroid/langroid/blob/main/langroid/language_models/prompt_formatter/hf_formatter.py)
that retrieves the HuggingFace tokenizer for the LLM and uses
its "apply_chat_template" method to format chats.
This gives you control over the chat format and you can use the "generation"
endpoint of the LLM API instead of the "chat" endpoint.

Once we switched to this, results improved dramatically 🚀

Be sure to checkout Langroid https://github.com/langroid/langroid

#llm #ai #opensource
</file>

<file path="docs/blog/posts/local-llm.md">
---
title: 'Using Langroid with Local LLMs'
draft: false
date: 2023-09-14
authors: 
  - pchalasani
categories:
  - langroid
  - llm
  - local-llm
comments: true
---
## Why local models?
There are commercial, remotely served models that currently appear to beat all open/local
models. So why care about local models? Local models are exciting for a number of reasons:

<!-- more -->

- **cost**: other than compute/electricity, there is no cost to use them.
- **privacy**: no concerns about sending your data to a remote server.
- **latency**: no network latency due to remote API calls, so faster response times, provided you can get fast enough inference.
- **uncensored**: some local models are not censored to avoid sensitive topics.
- **fine-tunable**: you can fine-tune them on private/recent data, which current commercial models don't have access to.
- **sheer thrill**: having a model running on your machine with no internet connection,
  and being able to have an intelligent conversation with it -- there is something almost magical about it.

The main appeal with local models is that with sufficiently careful prompting,
they may behave sufficiently well to be useful for specific tasks/domains,
and bring all of the above benefits. Some ideas on how you might use local LLMs:

- In a multi-agent system, you could have some agents use local models for narrow 
  tasks with a lower bar for accuracy (and fix responses with multiple tries).
- You could run many instances of the same or different models and combine their responses.
- Local LLMs can act as a privacy layer, to identify and handle sensitive data before passing to remote LLMs.
- Some local LLMs have intriguing features, for example llama.cpp lets you 
  constrain its output using a grammar.

## Running LLMs locally

There are several ways to use LLMs locally. See the [`r/LocalLLaMA`](https://www.reddit.com/r/LocalLLaMA/comments/11o6o3f/how_to_install_llama_8bit_and_4bit/) subreddit for
a wealth of information. There are open source libraries that offer front-ends
to run local models, for example [`oobabooga/text-generation-webui`](https://github.com/oobabooga/text-generation-webui)
(or "ooba-TGW" for short) but the focus in this tutorial is on spinning up a
server that mimics an OpenAI-like API, so that any code that works with
the OpenAI API (for say GPT3.5 or GPT4) will work with a local model,
with just a simple change: set `openai.api_base` to the URL where the local API
server is listening, typically `http://localhost:8000/v1`.

There are a few libraries we recommend for setting up local models with OpenAI-like APIs:

- [LiteLLM OpenAI Proxy Server](https://docs.litellm.ai/docs/proxy_server) lets you set up a local 
  proxy server for over 100+ LLM providers (remote and local).
- [ooba-TGW](https://github.com/oobabooga/text-generation-webui) mentioned above, for a variety of models, including llama2 models.
- [llama-cpp-python](https://github.com/abetlen/llama-cpp-python) (LCP for short), specifically for llama2 models.
- [ollama](https://github.com/jmorganca/ollama)

We recommend visiting these links to see how to install and run these libraries.

## Use the local model with the OpenAI library

Once you have a server running using any of the above methods, 
your code that works with the OpenAI models can be made to work 
with the local model, by simply changing the `openai.api_base` to the 
URL where the local server is listening. 

If you are using Langroid to build LLM applications, the framework takes
care of the `api_base` setting in most cases, and you need to only set
the `chat_model` parameter in the LLM config object for the LLM model you are using.
See the [Non-OpenAI LLM tutorial](../../tutorials/non-openai-llms.md) for more details.



<iframe src="https://langroid.substack.com/embed" width="480" height="320" style="border:1px solid #EEE; background:white;" frameborder="0" scrolling="no"></iframe>
</file>

<file path="docs/blog/posts/malade.md">
---
title: 'MALADE: Multi-Agent Architecture for Pharmacovigilance'
draft: false
date: 2024-08-12
authors:
- jihye
- nils
- pchalasani
- mengelhard
- someshjha
- anivaryakumar
- davidpage

categories:
- langroid
- multi-agent
- neo4j
- rag
comments: true
---

# MALADE: Multi-Agent Architecture for Pharmacovigilance

[Published in ML for HealthCare 2024](https://www.mlforhc.org/2024-abstracts)

[Arxiv](https://arxiv.org/abs/2408.01869) 

[GitHub](https://github.com/jihyechoi77/malade)

## Summary
We introduce MALADE (**M**ultiple **A**gents powered by **L**LMs for **ADE** Extraction),
a multi-agent system for Pharmacovigilance. It is the first effective explainable 
multi-agent LLM system for extracting Adverse Drug Events (ADEs) from FDA drug labels and drug prescription data.
<!-- more -->
Given a drug category and an adverse outcome, MALADE
produces:

- a qualitative label of risk (`increase`, `decrease` or `no-effect`),
- confidence in the label (a number in $[0,1]$),
- frequency of effect (`rare`, `common`, or `none`),
- strength of evidence (`none`, `weak`, or `strong`), and
- a justification with citations.

This task is challenging for several reasons: 

- FDA labels and prescriptions are for individual drugs, not drug categories, so representative drugs in a category 
  need to be identified from patient prescription data, and ADE information found for specific drugs in a category 
  needs to be aggregated to make a statement about the category as a whole, 
- The data is noisy, with variations in the terminologies of drugs and outcomes, and 
- ADE descriptions are often buried in large amounts of narrative text.

The MALADE architecture is LLM-agnostic 
and leverages the [Langroid](https://github.com/langroid/langroid) multi-agent framework.
It consists of a combination of Agents using Retrieval Augmented Generation (RAG), that 
iteratively improve their answers based on feedback from Critic Agents.
We evaluate the quantitative scores against 
a ground-truth dataset known as the [*OMOP Ground Truth Task*](https://www.niss.org/sites/default/files/Session3-DaveMadigan_PatrickRyanTalk_mar2015.pdf)
and find that MALADE achieves state-of-the-art performance.



## Introduction

In the era of Large Language Models (LLMs), given their remarkable text understanding and generation abilities, 
there is an unprecedented opportunity to develop new, LLM-based methods for trustworthy medical knowledge synthesis, 
extraction and summarization. The focus of this paper is Pharmacovigilance, a critical task in healthcare, where 
the goal is to monitor and evaluate the safety of drugs. In particular, the identification of Adverse Drug Events 
(ADEs) is crucial for ensuring patient safety. Consider a question such as this:

> What is the effect of **ACE inhibitors** on the risk of developing **angioedema**?

Here the **drug category** $C$ is _ACE inhibitors_, and the **outcome** $O$ is _angioedema_.
Answering this question involves several steps:

- **1(a): Find all drugs** in the ACE inhibitor category $C$, e.g. by searching the FDA 
[National Drug Code](https://www.fda.gov/drugs/drug-approvals-and-databases/national-drug-code-directory) (NDC) 
   database. This can be done using Elastic-Search, with filters to handle variations in drug/category names and inaccurate classifications.
- **1(b): Find the prescription frequency** of each drug in $C$ from patient prescription data, e.g. 
the [MIMIC-IV](https://physionet.org/content/mimiciv/3.0/) database. This can be done with a SQL query.
- **1(c): Identify the representative drugs** $D \subset C$ in this category, based on prescription frequency data 
     from step 2.  
- **2:** For each drug $d \in D$, **summarize ADE information** about the effect of $d$ on the outcome $O$ of interest,
   (in this case angioedema) from text-based pharmaceutical sources, 
    e.g. the [OpenFDA Drug Label](https://open.fda.gov/apis/drug/label/) database.
- **3: Aggregate** the information from all drugs in $D$ to make a statement about the category $C$ as a whole.


## The role of LLMs

While steps 1(a) and 1(b) can be done by straightforward deterministic algorithms (SQL queries or Elastic-Search), the 
remaining steps are challenging but ideally suited to LLMs:

### Step 1(c): Identifying representative drugs in a category from prescription frequency data (`DrugFinder` Agent)

This is complicated by noise, such as the same drug appearing multiple times under different names, 
formulations or delivery methods (For example, the ACE inhibitor **Lisinopril** is also known as **Zestril** and **Prinivil**.) 
  Thus a judgment must
  be made as to whether these are sufficiently different to be considered pharmacologically distinct;
  and some of these drugs may not actually belong to the category. This task thus requires a grouping operation, 
  related to the task of identifying standardized drug codes from text descriptions,
  well known to be challenging. This makes it very difficult to explicitly define the algorithm in a deterministic 
  manner that covers all edge cases (unlike the above database tasks), and hence is well-suited
  to LLMs, particularly those such as GPT-4, Claude3.5, and similar-strength variants which are known to have been 
  trained on vast amounts of general medical texts. 

In MALADE, this task is handled by the `DrugFinder` agent,
which is an Agent/Critic system where the main agent iteratively improves its output
in a feedback loop with the Critic agent. For example, the Critic corrects the Agent when it incorrectly
classifies drugs as pharmacologically distinct.

###  Step 2: Identifying Drug-Outcome Associations (`DrugOutcomeInfoAgent`)

The task here is to identify whether a given drug
has an established effect on the risk of a given outcome, based on FDA drug label database, and
output a summary of relevant information, including the level of identified risk and the evidence for
such an effect. Since this task involves extracting information from narrative text, it is well-suited to
LLMs using the Retrieval Augmented Generation (RAG) technique. 

In MALADE, the `DrugOutcomeInfoAgent` handles this task, and is also an Agent/Critic system, where the Critic
provides feedback and corrections to the Agent's output.
This agent does not have direct access to the FDA Drug Label data, but can receive
this information via another agent, `FDAHandler`. FDAHandler is equipped with **tools** (also known as function-calls) 
to invoke the OpenFDA API for drug label data, and answers questions in the context of information retrieved
based on the queries. Information received from this API is ingested into a vector database, so the
agent first uses a tool to query this vector database, and only resorts to the OpenFDA API tool if
the vector database does not contain the relevant information. An important aspect of this agent is that
its responses include specific **citations** and **excerpts** justifying its conclusions.

###  Step 3: Labeling Drug Category-Outcome Associations (`CategoryOutcomeRiskAgent`)

To identify association between a drug category C and an adverse health outcome $O$, we concurrently run a batch of 
queries to copies of `DrugOutcomeInfoAgent`, one for each drug $d$ in the
representative-list $D$ for the category, of the form: 

> Does drug $d$ increase or decrease the risk of condition $O$?

The results are sent to `CategoryOutcomeRiskAgent`, 
which is an Agent/Critic system which performs the final classification
step; its goal is to generate the qualitative and quantitative outputs mentioned above.

## MALADE Architecture

The figure below illustrates how the MALADE architecture handles the query,

> What is the effect of **ACE inhibitors** on the risk of developing **angioedema**?

![malade-arch.png](figures/malade-arch.png)

The query triggers a sequence of subtasks performed by the three Agents described above: 
`DrugFinder`, `DrugOutcomeInfoAgent`, and `CategoryOutcomeRiskAgent`.
Each Agent generates a response and justification, which are validated by a corresponding Critic agent, whose feedback is
used by the Agent to revise its response.

## Evaluation

### OMOP Ground Truth

We evaluate the results of MALADE against a well-established ground-truth dataset, 
the [OMOP ADE ground-truth table](https://www.niss.org/sites/default/files/Session3-DaveMadigan_PatrickRyanTalk_mar2015.pdf), shown below.
This is a reference dataset within the Observational Medical Outcomes Partnership (OMOP) Common Data Model that 
contains validated information about known adverse drug events.

![omop-ground-truth.png](figures/omop-ground-truth.png)

### Confusion Matrix

Below is a side-by-side comparison of this ground-truth dataset (left) with MALADE's labels (right), ignoring blue 
cells (see the paper for details):

![omop-results.png](figures/omop-results.png)

The resulting confusion-matrix for MALADE is shown below:

![confusion.png](figures/confusion.png)

### AUC Metric

Since MALADE produces qualitative and quantitative outputs, the paper explores a variety of ways to evaluate its
performance against the OMOP ground-truth dataset. Here we focus on the label output $L$ (i.e. `increase`, 
`decrease`, or `no-effect`), and its associated confidence score $c$, and use the Area Under the ROC Curve (AUC) as 
the evaluation metric.
The AUC metric is designed for binary classification, so we transform the three-class label output $L$ and
confidence score $c$ to a binary classification score $p$ as follows.
We treat $L$ = `increase` as the positive class,
and $L$ = `decrease` or `no-effect` as the negative class, and
we transform the label confidence score $c$ into a probability $p$ of `increase` as follows:


- if the label output is `increase`, $p = (2+c)/3$,
- if the label output is `no-effect`, $p = (2-c)/3$, and
- if the label output is `decrease` , $p = (1-c)/3$.

These transformations align with two intuitions: (a) a *higher* confidence in `increase` corresponds
to a *higher* probability of `increase`, and a *higher* confidence in `no-effect` or `decrease`
corresponds to a *lower* probability of `increase`, and (b) for a given confidence score $c$, the progression
of labels `decrease`, `no-effect`, and `increase` corresponds to *increasing* probabilities of `increase`.
The above transformations ensure that the probability $p$ is in the range $[0,1]$ and scales linearly with the
confidence score $c$.

We ran the full MALADE system for all drug-category/outcome pairs in the OMOP ground-truth dataset, 
and then computed the AUC for the score $p$ against the ground-truth binary classification label.
With `GPT-4-Turbo` we obtained an AUC of 0.85, while `GPT-4o` resulted in an AUC of 0.90.
These are state-of-the-art results for this specific ADE-extraction task.


### Ablations

An important question the paper investigates is whether (and how much) the various components (RAG, critic agents, etc)
contribute to MALADE's performance. To answer this, we perform ablations, where we remove one or more
components from the MALADE system and evaluate the performance of the resulting system.
For example we found that dropping the Critic agents reduces the AUC (using `GPT-4-Turbo`) from 0.85 to 0.82
(see paper, Appendix D for more ablation results).

### Variance of LLM-generated Scores

When using an LLM to generate numerical scores, it is important to understand the variance in the scores.
For example, if a single "full" run of MALADE (i.e. for all drug-category/outcome pairs in the OMOP ground-truth
dataset) produces a certain AUC, was it a "lucky" run, or is the AUC relatively stable across runs?
Ideally one would investigate this by repeating the full run of MALADE many times,  
but given the expense of running a full experiment, we focus on just three representative cells in the OMOP table,
one corresponding to each possible ground-truth label, and run MALADE 10 times for each cells, and
study the distribution of $p$ (the probability of increased risk, translated from the confidence score using the
method described above), for each output label. Encouragingly, we find that the distribution of $p$ shows clear
separation between the three labels, as in the figure below (The $x$ axis ranges from 0 to 1, and the three colored
groups of bars represent, from left to right, `decrease`, `no-effect`, and `increase` labels). Full details are in 
the Appendix D of the paper.

![img.png](figures/variance-histogram.png)
</file>

<file path="docs/blog/posts/multi-agent-debate.md">
---
title: 'Multi Agent Debate and Education Platform'
draft: false
date: 2025-02-04
authors: 
  - adamshams
categories:
  - langroid
  - llm
  - local-llm
  - chat
comments: true
---

## Introduction
Have you ever imagined a world where we can debate complex issues with Generative AI agents taking a distinct 
stance and backing their arguments with evidence? Some will change your mind, and some will reveal the societal biases 
on which each distinctive Large Language Model (LLM) is trained on. Introducing an [AI-powered debate platform](https://github.com/langroid/langroid/tree/main/examples/multi-agent-debate) that brings 
this imagination to reality, leveraging diverse LLMs and the Langroid multi-agent programming framework. The system enables users to engage in structured debates with an AI taking the opposite stance (or even two AIs debating each other), using a multi-agent architecture with Langroid's powerful framework, where each agent embodies a specific ethical perspective, creating realistic and dynamic interactions. 
Agents are prompt-engineered and role-tuned to align with their assigned ethical stance, 
ensuring thoughtful and structured debates. 

<!-- more -->

My motivations for creating this platform included: 

  - A debate coach for underserved students without access to traditional resources. 
  - Tool for research and generating arguments from authentic sources. 
  - Create an adaptable education platform to learn two sides of the coin for any topic.
  - Reduce echo chambers perpetuated by online algorithms by fostering two-sided debates on any topic, promoting education and awareness around misinformation. 
  - Provide a research tool to study the varieties of biases in LLMs that are often trained on text reflecting societal biases. 
  - Identify a good multi-agent framework designed for programming with LLMs.


## Platform Features:
### Dynamic Agent Generation:
The platform features five types of agents: Pro, Con, Feedback, Research, and Retrieval Augmented Generation (RAG) Q&A. 
Each agent is dynamically generated using role-tuned and engineered prompts, ensuring diverse and engaging interactions.
#### Pro and Con Agents: 
These agents engage in the core debate, arguing for and against the chosen topic. 
Their prompts are carefully engineered to ensure they stay true to their assigned ethical stance.
#### Feedback Agent: 
This agent provides real-time feedback on the arguments and declares a winner. The evaluation criteria are based on the well-known Lincoln–Douglas debate format, and include:

  - Clash of Values 
  - Argumentation 
  - Cross-Examination 
  - Rebuttals 
  - Persuasion 
  - Technical Execution 
  - Adherence to Debate Etiquette 
  - Final Focus
#### Research Agent: 
This agent has the following functionalities:

  - Utilizes the `MetaphorSearchTool` and the `Metaphor` (now called `Exa`) Search API to conduct web searches combined with
Retrieval Augmented Generation (RAG) to relevant web references for user education about the selected topic. 
  - Produces a summary of arguments for and against the topic.
  - RAG-based document chat with the resources identified through Web Search. 
#### RAG Q&A Agent:

  - Provides Q&A capability using a RAG based chat interaction with the resources identified through Web Search.
The agent utilizes `DocChatAgent` that is part of Langroid framework which orchestrates all LLM interactions. 
  - Rich chunking parameters allows the user to get optimized relevance results. Check out `config.py`for details.

### Topic Adaptability:
Easily adaptable to any subject by simply adding pro and con system messages. This makes it a versatile tool for
exploring diverse topics and fostering critical thinking. Default topics cover ethics and use of AI for the following:
  - Healthcare
  - Intellectual property 
  - Societal biases 
  - Education
### Autonomous or Interactive:
Engage in manual debate with a pro or con agent or watch it autonomously while adjusting number of turns.

### Diverse LLM Selection Adaptable per Agent: 
Configurable to select from diverse commercial and open source models: OpenAI, Google, and Mistral 
to experiment with responses for diverse perspectives. Users can select a unique LLM for each agent. 
       
### LLM Tool/Function Integration: 
Utilizes LLM tools/functions features to conduct semantic search using Metaphor Search API and summarizes the pro and 
con perspectives for education.

### Configurable LLM Parameters: 
Parameters like temperature, minimum and maximum output tokens, allowing for customization of the AI's responses.
Configurable LLM parameters like temperature, min & max output tokens. For Q&A with the searched resources, several
parameters can be tuned in the `config` to enhance response relevance.

### Modular Design: 
Reusable code and modularized for other LLM applications.


## Interaction
1. Decide if you want to you use same LLM for all agents or different ones
2. Decide if you want autonomous debate between AI Agents or user vs. AI Agent. 
3. Select a debate topic.
4. Choose your side (Pro or Con).
5. Engage in a debate by providing arguments and receiving responses from agents.
6. Request feedback at any time by typing `f`.
7. Decide if you want the Metaphor Search to run to find Topic relevant web links
   and summarize them. 
8. Decide if you want to chat with the documents extracted from URLs found to learn more about the Topic.
9. End the debate manually by typing `done`. If you decide to chat with the documents, you can end session
by typing `x`

## Why was Langroid chosen?
I chose Langroid framework because it's a principled multi-agent programming framework inspired by the Actor framework.
Prior to using Langroid, I developed a multi-agent debate system, however, I had to write a lot of tedious code to manage states of communication between
debating agents, and the user interactions with LLMs. Langroid allowed me to seamlessly integrate multiple LLMs,
easily create agents, tasks, and attach sub-tasks. 

### Agent Creation Code Example

```python
   def create_chat_agent(name: str, llm_config: OpenAIGPTConfig, system_message: str) -> ChatAgent:
   
    return ChatAgent(
        ChatAgentConfig(
            llm=llm_config,
            name=name,
            system_message=system_message,
        )
    )
```
#### Sample Pro Topic Agent Creation

```python
 
    pro_agent = create_chat_agent(
        "Pro",
        pro_agent_config,
        system_messages.messages[pro_key].message + DEFAULT_SYSTEM_MESSAGE_ADDITION,
    )
    
```
The `Task` mechanism in Langroid provides a robust mechanism for managing complex interactions within multi-agent 
systems. `Task` serves as a container for managing the flow of interactions between different agents
(such as chat agents) and attached sub-tasks.`Task` also helps with turn-taking, handling responses, 
and ensuring smooth transitions between dialogue states. Each Task object is responsible for coordinating responses 
from its assigned agent, deciding the sequence of responder methods (llm_response, user_response, agent_response), 
and managing transitions between different stages of a conversation or debate. Each agent can focus on its specific 
role while the task structure handles the overall process's orchestration and flow, allowing a clear separation of 
concerns. The architecture and code transparency of Langroid's framework make it an incredible candidate for 
applications like debates where multiple agents must interact dynamically and responsively
based on a mixture of user inputs and automated responses.

### Task creation and Orchestration Example

```python
    user_task = Task(user_agent, interactive=interactive_setting, restart=False)
    ai_task = Task(ai_agent, interactive=False, single_round=True)
    user_task.add_sub_task(ai_task)
    if not llm_delegate:
        user_task.run(user_agent.user_message, turns=max_turns)
    else:
        user_task.run("get started", turns=max_turns)
    
```
Tasks can be easily set up as sub-tasks of an orchestrating agent. In this case user_task could be Pro or Con depending 
on the user selection. 

If you want to build custom tools/functions or use Langroid provided it is only a line of code using
`agent.enable_messaage`. Here is an example of `MetaphorSearchTool` and `DoneTool`. 
```python
        metaphor_search_agent.enable_message(MetaphorSearchTool)
        metaphor_search_agent.enable_message(DoneTool)
```

Overall I had a great learning experience using Langroid and recommend using it for any projects 
that need to utilize LLMs. I am already working on a few Langroid based information retrieval and research systems 
for use in medicine and hoping to contribute more soon. 

### Bio

I'm a high school senior at Khan Lab School located in Mountain View, CA where I host a student-run Podcast known as the
Khan-Cast. I also enjoy tinkering with interdisciplinary STEM projects. You can reach me on [LinkedIn](https://www.linkedin.com/in/adamshams/).
</file>

<file path="docs/blog/posts/test.md">
---
draft: true
date: 2022-01-31
authors: 
  - pchalasani
categories:
  - test
  - blog
comments: true
---

# Test code snippets

```python
from langroid.language_models.base import LLMMessage, Role
msg = LLMMessage(
        content="What is the capital of Bangladesh?",
        role=Role.USER,
      )
```

<!-- more -->


# Test math notation

A nice equation is $e^{i\pi} + 1 = 0$, which is known as Euler's identity.
Here is a cool equation too, and in display mode:

$$
e = mc^2
$$

# Latex with newlines

Serious latex with `\\` for newlines renders fine:

$$
\begin{bmatrix}
a & b \\
c & d \\
e & f \\
\end{bmatrix}
$$

or a multi-line equation

$$
\begin{aligned}
\dot{x} & = \sigma(y-x) \\
\dot{y} & = \rho x - y - xz \\
\dot{z} & = -\beta z + xy
\end{aligned}
$$

<iframe src="https://langroid.substack.com/embed" width="480" height="320" style="border:1px solid #EEE; background:white;" frameborder="0" scrolling="no"></iframe>
</file>

<file path="docs/blog/.authors.yml">
authors:
  pchalasani:
    name: Prasad Chalasani
    description: Langroid CoFounder
    avatar: https://github.com/pchalasani.png
  mohannad:
    name: Mohannad Alhanahnah
    description: Langroid Contributor
    avatar: https://avatars.githubusercontent.com/u/15859139
  nils:
    name: Nils Palumbo
    description: Phd Candidate (CS), UW-Madison; Langroid core dev.
    avatar: https://www.github.com/nilspalumbo.png
  jihye:
    name: Jihye Choi
    description: PhD Candidate (CS), UW-Madison
    avatar: https://www.github.com/jihyechoi77.png
  someshjha:
    name: Somesh Jha
    description: UW-Madison; Langroid CoFounder
    avatar: https://www.gravatar.com/avatar/?d=mp
  anivaryakumar:
    name: Anivarya Kumar
    description: Duke University
    avatar: https://www.gravatar.com/avatar/?d=mp
  davidpage:
    name: David Page
    description: Duke University
    avatar: https://www.gravatar.com/avatar/?d=mp
  mengelhard:
    name: Matthew Engelhard
    description: Duke University
    avatar: https://www.gravatar.com/avatar/?d=mp
  adamshams:
    name: Adam Shams
    description: Langroid Contributor, Khan Lab School
    avatar: https://avatars.githubusercontent.com/u/84205479
</file>

<file path="docs/blog/index.md">
# Blog
</file>

<file path="docs/demos/targeting/audience-targeting.md">
# Audience Targeting for a Business

Suppose you are a marketer for a business, trying to figure out which 
audience segments to target.
Your downstream systems require that you specify _standardized_ audience segments
to target, for example from the [IAB Audience Taxonomy](https://iabtechlab.com/standards/audience-taxonomy/).

There are thousands of standard audience segments, and normally you would need 
to search the list for potential segments that match what you think your ideal
customer profile is. This is a tedious, error-prone task.

But what if we can leverage an LLM such as GPT-4?
We know that GPT-4 has  skills that are ideally suited for this task:

- General knowledge about businesses and their ideal customers
- Ability to recognize which standard segments match an English description of a customer profile
- Ability to plan a conversation to get the information it needs to answer a question


Once you decide to use an LLM, you still need to figure out how to organize the 
various components of this task:

- **Research:** What are some ideal customer profiles for the business
- **Segmentation:** Which standard segments match an English description of a customer profile
- **Planning:** how to organize the task to identify a few standard segments

## Using Langroid Agents 

Langroid makes it intuitive and simple to build an LLM-powered system organized
around agents, each responsible for a different task.
In less than a day we built a 3-agent system to automate this task:

- The `Marketer` Agent is given the Planning role.
- The `Researcher` Agent is given the Research role, 
  and it has access to the business description. 
- The `Segmentor` Agent is given the Segmentation role. It has access to the 
  IAB Audience Taxonomy via a vector database, i.e. its rows have been mapped to
  vectors via an embedding model, and these vectors are stored in a vector-database. 
  Thus given an English description of a customer profile,
  the `Segmentor` Agent maps it to a vector using the embedding model,
  and retrieves the nearest (in vector terms, e.g. cosine similarity) 
  IAB Standard Segments from the vector-database. The Segmentor's LLM 
  further refines this by selecting the best-matching segments from the retrieved list.

To kick off the system, the human user describes a business in English,
or provides the URL of the business's website. 
The `Marketer` Agent sends
customer profile queries to the `Researcher`, who answers in plain English based on 
the business description, and the Marketer takes this description and sends it to the Segmentor,
who maps it to Standard IAB Segments. The task is done when the Marketer finds 4 Standard segments. 
The agents are depicted in the diagram below:

![targeting.png](targeting.png)

## An example: Glashutte Watches

The human user first provides the URL of the business, in this case:
```text
https://www.jomashop.com/glashutte-watches.html
```
From this URL, the `Researcher` agent summarizes its understanding of the business.
The `Marketer` agent starts by asking the `Researcher`:
``` 
Could you please describe the age groups and interests of our typical customer?
```
The `Researcher` responds with an English description of the customer profile:
```text
Our typical customer is a fashion-conscious individual between 20 and 45 years...
```
The `Researcher` forwards this English description to the `Segmentor` agent, who
maps it to a standardized segment, e.g.:
```text
Interest|Style & Fashion|Fashion Trends
...
```
This conversation continues until the `Marketer` agent has identified 4 standardized segments.

Here is what the conversation looks like:

![targeting.gif](targeting.gif)
</file>

<file path="docs/examples/agent-tree.md">
# Hierarchical computation with Langroid Agents 

Here is a simple example showing tree-structured computation
where each node in the tree is handled by a separate agent.
This is a toy numerical example, and illustrates:

- how to have agents organized in a hierarchical structure to accomplish a task 
- the use of global state accessible to all agents, and 
- the use of tools/function-calling.

## The Computation 

We want to carry out the following calculation for a given input number $n$:

```python
def Main(n):
    if n is odd:
        return (3*n+1) + n
    else:
        if n is divisible by 10:
            return n/10 + n
        else:
            return n/2 + n
```

## Using function composition

Imagine we want to do this calculation using a few auxiliary functions:

```python
def Main(n):
    # return non-null value computed by Odd or Even
    Record n as global variable # to be used by Adder below
    return Odd(n) or Even(n)

def Odd(n):
    # Handle odd n
    if n is odd:
        new = 3*n+1
        return Adder(new)
    else:
        return None
    
def Even(n):
    # Handle even n: return non-null value computed by EvenZ or EvenNZ
    return EvenZ(n) or EvenNZ(n)

def EvenZ(n):
    # Handle even n divisible by 10, i.e. ending in Zero
    if n is divisible by 10:
        new = n/10
        return Adder(new)
    else:
        return None
    
def EvenNZ(n):
    # Handle even n not divisible by 10, i.e. not ending in Zero
    if n is not divisible by 10:
        new = n/2
        return Adder(new)
    else:
        return None  

def Adder(new):
    # Add new to starting number, available as global variable n
    return new + n
```

## Mapping to a tree structure

This compositional/nested computation can be represented as a tree:

```plaintext
       Main
     /     \
  Even     Odd
  /   \        \
EvenZ  EvenNZ   Adder
  |      |
 Adder  Adder
```

Let us specify the behavior we would like for each node, in a 
"decoupled" way, i.e. we don't want a node to be aware of the other nodes.
As we see later, this decoupled design maps very well onto Langroid's
multi-agent task orchestration. To completely define the node behavior,
we need to specify how it handles an "incoming" number $n$ (from a parent node 
or user), and how it handles a "result" number $r$ (from a child node).

- `Main`: 
    - incoming $n$: simply send down $n$, record the starting number $n_0 = n$ as a global variable. 
    - result $r$: return $r$.
- `Odd`: 
    - incoming $n$: if n is odd, send down $3*n+1$, else return None
    - result $r$: return $r$
- `Even`: 
    - incoming $n$: if n is even, send down $n$, else return None
    - result $r$: return $r$
- `EvenZ`: (guaranteed by the tree hierarchy, to receive an even number.)  
    - incoming $n$: if n is divisible by 10, send down $n/10$, else return None
    - result $r$: return $r$
- `EvenNZ`: (guaranteed by the tree hierarchy, to receive an even number.)
    - incoming $n$: if n is not divisible by 10, send down $n/2$, else return None
    - result $r$: return $r$
- `Adder`:
    - incoming $n$: return $n + n_0$ where $n_0$ is the 
    starting number recorded by Main as a global variable.
    - result $r$: Not applicable since `Adder` is a leaf node.
  
## From tree nodes to Langroid Agents 

Let us see how we can perform this calculation using multiple Langroid agents, where

- we define an agent corresponding to each of the nodes above, namely 
`Main`, `Odd`, `Even`, `EvenZ`, `EvenNZ`, and `Adder`.
- we wrap each Agent into a Task, and use the `Task.add_subtask()` method to connect the agents into 
  the desired hierarchical structure.

Below is one way to do this using Langroid. We designed this with the following
desirable features:

- Decoupling: Each agent is instructed separately, without mention of any other agents
  (E.g. Even agent does not know about Odd Agent, EvenZ agent, etc).
  In particular, this means agents will not be "addressing" their message
  to specific other agents, e.g. send number to Odd agent when number is odd,
  etc. Allowing addressing would make the solution easier to implement,
  but would not be a decoupled solution.
  Instead, we want Agents to simply put the number "out there", and have it handled
  by an applicable agent, in the task loop (which consists of the agent's responders,
  plus any sub-task `run` methods).

- Simplicity: Keep the agent instructions relatively simple. We would not want a solution
  where we have to instruct the agents (their LLMs) in convoluted ways. 

One way naive solutions fail is because agents are not able to distinguish between
a number that is being "sent down" the tree as input, and a number that is being
"sent up" the tree as a result from a child node.

We use a simple trick: we instruct the LLM to mark returned values using the RESULT keyword,
and instruct the LLMs on how to handle numbers that come with RESULT keyword, and those that don't
In addition, we leverage some features of Langroid's task orchestration:

- When `llm_delegate` is `True`, if the LLM says `DONE [rest of msg]`, the task is
  considered done, and the result of the task is `[rest of msg]` (i.e the part after `DONE`).
- In the task loop's `step()` function (which seeks a valid message during a turn of
  the conversation) when any responder says `DO-NOT-KNOW`, it is not considered a valid
  message, and the search continues to other responders, in round-robin fashion.



See the [`chat-tree.py`](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tree.py)
example for an implementation of this solution. You can run that example as follows:
```bash
python3 examples/basic/chat-tree.py
```
In the sections below we explain the code in more detail.

## Define the agents

Let us start with defining the configuration to be used by all agents:

```python
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig

config = ChatAgentConfig(
  llm=OpenAIGPTConfig(
    chat_model=OpenAIChatModel.GPT4o,
  ),
  vecdb=None, # no need for a vector database
)
```

Next we define each of the agents, for example:

```python
main_agent = ChatAgent(config)
```

and similarly for the other agents.

## Wrap each Agent in a Task

To allow agent interactions, the first step is to wrap each agent in a Task.
When we define the task, we pass in the instructions above as part of the system message.
Recall the instructions for the `Main` agent:

- `Main`:
    - incoming $n$: simply send down $n$, record the starting number $n_0 = n$ as a global variable.
    - result $r$: return $r$.

We include the equivalent of these instructions in the `main_task` that wraps 
the `main_agent`:

```python
from langroid.agent.task import Task

main_task = Task(
    main_agent,
    name="Main",
    interactive=False, #(1)!
    system_message="""
          You will receive two types of messages, to which you will respond as follows:
          
          INPUT Message format: <number>
          In this case simply write the <number>, say nothing else.
          
          RESULT Message format: RESULT <number>
          In this case simply say "DONE <number>", e.g.:
          DONE 19
    
          To start off, ask the user for the initial number, 
          using the `ask_num` tool/function.
          """,
    llm_delegate=True, # allow LLM to control end of task via DONE
    single_round=False,
)
```

1. Non-interactive: don't wait for user input in each turn 

There are a couple of points to highlight about the `system_message` 
value in this task definition:

- When the `Main` agent receives just a number, it simply writes out that number,
  and in the Langroid Task loop, this number becomes the "current pending message"
  to be handled by one of the sub-tasks, i.e. `Even, Odd`. Note that these sub-tasks
  are _not_ mentioned in the system message, consistent with the decoupling principle.
- As soon as either of these sub-tasks returns a non-Null response, in the format "RESULT <number>", the `Main` agent
  is instructed to return this result saying "DONE <number>". Since `llm_delegate`
  is set to `True` (meaning the LLM can decide when the task has ended), 
  this causes the `Main` task to be considered finished and the task loop is exited.

Since we want the `Main` agent to record the initial number as a global variable,
we use a tool/function `AskNum` defined as follows 
(see [this section](../quick-start/chat-agent-tool.md) in the getting started guide 
for more details on Tools):

```python
from rich.prompt import Prompt
from langroid.agent.tool_message import ToolMessage


class AskNumTool(ToolMessage):
  request = "ask_num"
  purpose = "Ask user for the initial number"

  def handle(self) -> str:
    """
    This is a stateless tool (i.e. does not use any Agent member vars), so we can
    define the handler right here, instead of defining an `ask_num`
    method in the agent.
    """
    num = Prompt.ask("Enter a number")
    # record this in global state, so other agents can access it
    MyGlobalState.set_values(number=num)
    return str(num)
```

We then enable the `main_agent` to use and handle messages that conform to the 
`AskNum` tool spec:

```python
main_agent.enable_message(AskNumTool)
```

!!! tip "Using and Handling a tool/function"
    "Using" a tool means the agent's LLM _generates_ 
    the function-call (if using OpenAI function-calling) or 
    the JSON structure (if using Langroid's native tools mechanism) 
    corresponding to this tool. "Handling" a tool refers to the Agent's method 
    recognizing the tool and executing the corresponding code.


The tasks for other agents are defined similarly. We will only note here
that the `Adder` agent needs a special tool `AddNumTool` to be able to add the current number
to the initial number set by the `Main` agent. 

## Connect the tasks into a tree structure

So far, we have wrapped each agent in a task, in isolation, and there is no 
connection between the tasks. The final step is to connect the tasks to 
the tree structure we saw earlier:

```python
main_task.add_sub_task([even_task, odd_task])
even_task.add_sub_task([evenz_task, even_nz_task])
evenz_task.add_sub_task(adder_task)
even_nz_task.add_sub_task(adder_task)
odd_task.add_sub_task(adder_task)
```

Now all that remains is to run the main task:

```python
main_task.run()
```

Here is what a run starting with $n=12$ looks like:

![chat-tree.png](chat-tree.png)
</file>

<file path="docs/javascripts/mathjax.js">
window.MathJax = {
  tex: {
    inlineMath: [["\\(", "\\)"]],
    displayMath: [["\\[", "\\]"]],
    processEscapes: true,
    processEnvironments: true
  },
  options: {
    ignoreHtmlClass: ".*|",
    processHtmlClass: "arithmatex"
  }
};

document$.subscribe(() => { 
  MathJax.typesetPromise()
})
</file>

<file path="docs/notes/async-streaming.md">
# Suppressing output in async, streaming mode

Available since version 0.18.0

When using an LLM API in streaming + async mode, you may want to suppress output,
especially when concurrently running multiple instances of the API.
To suppress output in async + stream mode, 
you can set the `async_stream_quiet` flag in [`LLMConfig`][langroid.language_models.base.LLMConfig]
to `True` (this is the default). 
Note that [`OpenAIGPTConfig`][langroid.language_models.openai_gpt.OpenAIGPTConfig]
inherits from `LLMConfig`, so you can use this flag with `OpenAIGPTConfig` as well:

```python
import langroid.language_models as lm
llm_config = lm.OpenAIGPTConfig(
    async_stream_quiet=True,
    ...
)
```
</file>

<file path="docs/notes/azure-openai-models.md">
# Azure OpenAI Models

To use OpenAI models deployed on Azure, first ensure a few environment variables
are defined (either in your `.env` file or in your environment):

- `AZURE_OPENAI_API_KEY`, from the value of `API_KEY`
- `AZURE_OPENAI_API_BASE` from the value of `ENDPOINT`, typically looks like `https://your_resource.openai.azure.com`.
- For `AZURE_OPENAI_API_VERSION`, you can use the default value in `.env-template`, and latest version can be found [here](https://learn.microsoft.com/en-us/azure/ai-services/openai/whats-new#azure-openai-chat-completion-general-availability-ga)
- `AZURE_OPENAI_DEPLOYMENT_NAME` is an OPTIONAL deployment name which may be
  defined by the user during the model setup.
- `AZURE_OPENAI_CHAT_MODEL` Azure OpenAI allows specific model names when you select the model for your deployment. You need to put precisely the exact model name that was selected. For example, GPT-3.5 (should be `gpt-35-turbo-16k` or `gpt-35-turbo`) or GPT-4 (should be `gpt-4-32k` or `gpt-4`).
- `AZURE_OPENAI_MODEL_NAME` (Deprecated, use `AZURE_OPENAI_CHAT_MODEL` instead).

This page [Microsoft Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/chatgpt-quickstart?tabs=command-line&pivots=programming-language-python#environment-variables) 
provides more information on how to obtain these values.

To use an Azure-deployed model in Langroid, you can use the `AzureConfig` class:

```python
import langroid.language_models as lm
import langroid as lr

llm_config = lm.AzureConfig(
    chat_model="gpt-4o"
    # the other settings can be provided explicitly here, 
    # or are obtained from the environment
)
llm = lm.AzureGPT(config=llm_config)

response = llm.chat(
  messages=[
    lm.LLMMessage(role=lm.Role.SYSTEM, content="You are a helpful assistant."),
    lm.LLMMessage(role=lm.Role.USER, content="3+4=?"),
  ]
)

agent = lr.ChatAgent(
    lr.ChatAgentConfig(
        llm=llm_config,
        system_message="You are a helpful assistant.",
    )
)

response = agent.llm_response("is 4 odd?")
print(response.content)  # "Yes, 4 is an even number."
response = agent.llm_response("what about 2?")  # follow-up question
```

## Using Azure OpenAI API v1 with Standard OpenAI Clients

Azure's October 2025 API update allows using standard OpenAI clients instead of
Azure-specific ones. However, Azure deployment names often differ from actual
model identifiers, which can cause issues with model capability detection.

If your deployment name differs from the actual model name, use `chat_model_orig`
to specify the actual model for proper capability detection:

```python
import langroid.language_models as lm

llm_config = lm.OpenAIGPTConfig(
    chat_model="my-gpt4o-deployment",     # Your Azure deployment name
    chat_model_orig="gpt-4o",             # Actual model name for capability detection
    api_base="https://your-resource.openai.azure.com/",
)
```

This ensures Langroid correctly identifies model capabilities (context length,
supported features, etc.) even when the deployment name doesn't match the
underlying model.
</file>

<file path="docs/notes/chunking.md">
# Document Chunking/Splitting in Langroid

Langroid's [`ParsingConfig`][langroid.parsing.parser.ParsingConfig]
provides several document chunking strategies through the `Splitter` enum:

## 1. MARKDOWN (`Splitter.MARKDOWN`) (The default)

**Purpose**: Structure-aware splitting that preserves markdown formatting.

**How it works**:

- Preserves document hierarchy (headers and sections)
- Enriches chunks with header information
- Uses word count instead of token count (with adjustment factor)
- Supports "rollup" to maintain document structure
- Ideal for markdown documents where preserving formatting is important

## 2. TOKENS (`Splitter.TOKENS`)

**Purpose**: Creates chunks of approximately equal token size.

**How it works**:

- Tokenizes the text using tiktoken
- Aims for chunks of size `chunk_size` tokens (default: 200)
- Looks for natural breakpoints like punctuation or newlines
- Prefers splitting at sentence/paragraph boundaries
- Ensures chunks are at least `min_chunk_chars` long (default: 350)

## 3. PARA_SENTENCE (`Splitter.PARA_SENTENCE`)

**Purpose**: Splits documents respecting paragraph and sentence boundaries.

**How it works**:

- Recursively splits documents until chunks are below 1.3× the target size
- Maintains document structure by preserving natural paragraph breaks
- Adjusts chunk boundaries to avoid cutting in the middle of sentences
- Stops when it can't split chunks further without breaking coherence

## 4. SIMPLE (`Splitter.SIMPLE`)

**Purpose**: Basic splitting using predefined separators.

**How it works**:

- Uses a list of separators to split text (default: `["\n\n", "\n", " ", ""]`)
- Splits on the first separator in the list
- Doesn't attempt to balance chunk sizes
- Simplest and fastest splitting method


## Basic Configuration

```python
from langroid.parsing.parser import ParsingConfig, Splitter

config = ParsingConfig(
    splitter=Splitter.MARKDOWN,  # Most feature-rich option
    chunk_size=200,              # Target tokens per chunk
    chunk_size_variation=0.30,   # Allowed variation from target
    overlap=50,                  # Token overlap between chunks
    token_encoding_model="text-embedding-3-small"
)
```

## Format-Specific Configuration

```python
# Customize PDF parsing
config = ParsingConfig(
    splitter=Splitter.PARA_SENTENCE,
    pdf=PdfParsingConfig(
        library="pymupdf4llm"  # Default PDF parser
    )
)

# Use Gemini for PDF parsing
config = ParsingConfig(
    pdf=PdfParsingConfig(
        library="gemini",
        gemini_config=GeminiConfig(
            model_name="gemini-2.0-flash",
            requests_per_minute=5
        )
    )
)
```

# Setting Up Parsing Config in DocChatAgentConfig

You can configure document parsing when creating a `DocChatAgent` by customizing the `parsing` field within the `DocChatAgentConfig`. Here's how to do it:

```python
from langroid.agent.special.doc_chat_agent import DocChatAgentConfig  
from langroid.parsing.parser import ParsingConfig, Splitter, PdfParsingConfig

# Create a DocChatAgent with custom parsing configuration
agent_config = DocChatAgentConfig(
    parsing=ParsingConfig(
        # Choose the splitting strategy
        splitter=Splitter.MARKDOWN,  # Structure-aware splitting with header context
        
        # Configure chunk sizes
        chunk_size=800,              # Target tokens per chunk
        overlap=150,                 # Overlap between chunks
        
        # Configure chunk behavior
        max_chunks=5000,             # Maximum number of chunks to create
        min_chunk_chars=250,         # Minimum characters when truncating at punctuation
        discard_chunk_chars=10,      # Discard chunks smaller than this
        
        # Configure context window
        n_neighbor_ids=3,            # Store 3 chunk IDs on either side
        
        # Configure PDF parsing specifically
        pdf=PdfParsingConfig(
            library="pymupdf4llm",   # Choose PDF parsing library
        )
    )
)
```
</file>

<file path="docs/notes/crawl4ai.md">
# Crawl4ai Crawler Documentation

## Overview

The `Crawl4aiCrawler` is a highly advanced and flexible web crawler integrated into Langroid, built on the powerful `crawl4ai` library. It uses a real browser engine (Playwright) to render web pages, making it exceptionally effective at handling modern, JavaScript-heavy websites. This crawler provides a rich set of features for simple page scraping, deep-site crawling, and sophisticated data extraction, making it the most powerful crawling option available in Langroid.

It is a local crawler, so no need for API keys.

## Installation

To use `Crawl4aiCrawler`, you must install the `crawl4ai` extra dependencies.

To install and prepare crawl4ai:

```bash
# Install langroid with crawl4ai support
pip install "langroid[crawl4ai]"
crawl4ai setup
crawl4ai doctor

```

> **Note**: The `crawl4ai setup` command will download Playwright browsers (Chromium, Firefox, WebKit) on first run. This is a one-time download that can be several hundred MB in size. The browsers are stored locally and used for rendering web pages.

## Key Features

- **Real Browser Rendering**: Accurately processes dynamic content, single-page applications (SPAs), and sites that require JavaScript execution.

- **Simple and Deep Crawling**: Can scrape a list of individual URLs (`simple` mode) or perform a recursive, deep crawl of a website starting from a seed URL (`deep` mode).

- **Powerful Extraction Strategies**:

  - **Structured JSON (No LLM)**: Extract data into a predefined JSON structure using CSS selectors, XPath, or Regex patterns. This is extremely fast, reliable, and cost-effective.

  - **LLM-Based Extraction**: Leverage Large Language Models (like GPT or Gemini) to extract data from unstructured content based on natural language instructions and a Pydantic schema.

- **Advanced Markdown Generation**: Go beyond basic HTML-to-markdown conversion. Apply content filters to prune irrelevant sections (sidebars, ads, footers) or use an LLM to intelligently reformat content for maximum relevance, perfect for RAG pipelines.

- **High-Performance Scraping**: Optionally use an LXML-based scraping strategy for a significant speed boost on large HTML documents.

- **Fine-Grained Configuration**: Offers detailed control over browser behavior (`BrowserConfig`) and individual crawl runs (`CrawlerRunConfig`) for advanced use cases.

## Configuration (`Crawl4aiConfig`)

The `Crawl4aiCrawler` is configured via the `Crawl4aiConfig` object. This class acts as a high-level interface to the underlying `crawl4ai` library's settings.

All of the strategies are optional.
Learn more about these strategies , browser_config and run_config at [Crawl4AI docs](https://docs.crawl4ai.com/)

```python
from langroid.parsing.url_loader import Crawl4aiConfig

# All parameters are optional and have sensible defaults
config = Crawl4aiConfig(
    crawl_mode="simple",  # or "deep"
    extraction_strategy=...,
    markdown_strategy=...,
    deep_crawl_strategy=...,
    scraping_strategy=...,
    browser_config=...,  # For advanced browser settings
    run_config=...,      # For advanced crawl-run settings
)
```

**Main Parameters:**

- `crawl_mode` (str):

  - `"simple"` (default): Crawls each URL in the provided list individually.

  - `"deep"`: Starts from the first URL in the list and recursively crawls linked pages based on the `deep_crawl_strategy`.

  - Make sure you are setting `"crawl_mode=deep"` whenever you are deep crawling this is crucial for smooth functioning.

- `extraction_strategy` (`ExtractionStrategy`): Defines how to extract structured data from a page. If set, the `Document.content` will be a **JSON string** containing the extracted data.

- `markdown_strategy` (`MarkdownGenerationStrategy`): Defines how to convert HTML to markdown. This is used when `extraction_strategy` is not set. The `Document.content` will be a **markdown string**.

- `deep_crawl_strategy` (`DeepCrawlStrategy`): Configuration for deep crawling, such as `max_depth`, `max_pages`, and URL filters. Only used when `crawl_mode` is `"deep"`.

- `scraping_strategy` (`ContentScrapingStrategy`): Specifies the underlying HTML parsing engine. Useful for performance tuning.

- `browser_config` & `run_config`: For advanced users to pass detailed `BrowserConfig` and `CrawlerRunConfig` objects directly from the `crawl4ai` library.

---

## Usage Examples

These are representative examples. For runnable examples check the script [`examples/docqa/crawl4ai_examples.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/crawl4ai_examples.py)

### 1. Simple Crawling (Default Markdown)

This is the most basic usage. It will fetch the content of each URL and convert it to clean markdown.

```python
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig

urls = [
    "https://pytorch.org/",
    "https://techcrunch.com/",
]

# Use default settings
crawler_config = Crawl4aiConfig()
loader = URLLoader(urls=urls, crawler_config=crawler_config)

docs = loader.load()
for doc in docs:
    print(f"URL: {doc.metadata.source}")
    print(f"Content (first 200 chars): {doc.content[:200]}")
```

### 2. Structured JSON Extraction (No LLM)

When you need to extract specific, repeated data fields from a page, schema-based extraction is the best choice. It's fast, precise, and free of LLM costs. The result in `Document.content` is a JSON string.

#### a. Using CSS Selectors (`JsonCssExtractionStrategy`)

This example scrapes titles and links from the Hacker News front page.

```python
import json
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.extraction_strategy import JsonCssExtractionStrategy

HACKER_NEWS_URL = "https://news.ycombinator.com"
HACKER_NEWS_SCHEMA = {
    "name": "HackerNewsArticles",
    "baseSelector": "tr.athing",
    "fields": [
        {"name": "title", "selector": "span.titleline > a", "type": "text"},
        {"name": "link", "selector": "span.titleline > a", "type": "attribute", "attribute": "href"},
    ],
}

# Create the strategy and pass it to the config
css_strategy = JsonCssExtractionStrategy(schema=HACKER_NEWS_SCHEMA)
crawler_config = Crawl4aiConfig(extraction_strategy=css_strategy)

loader = URLLoader(urls=[HACKER_NEWS_URL], crawler_config=crawler_config)
documents = loader.load()

# The Document.content will contain the JSON string
extracted_data = json.loads(documents[0].content)
print(json.dumps(extracted_data[:3], indent=2))
```

#### b. Using Regex (`RegexExtractionStrategy`)

This is ideal for finding common patterns like emails, URLs, or phone numbers.

```python
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.extraction_strategy import RegexExtractionStrategy

url = "https://www.scrapethissite.com/pages/forms/"

# Combine multiple built-in patterns
regex_strategy = RegexExtractionStrategy(
    pattern=(
        RegexExtractionStrategy.Email
        | RegexExtractionStrategy.Url
        | RegexExtractionStrategy.PhoneUS
    )
)

crawler_config = Crawl4aiConfig(extraction_strategy=regex_strategy)
loader = URLLoader(urls=[url], crawler_config=crawler_config)
documents = loader.load()

print(documents[0].content)
```

### 3. Advanced Markdown Generation

For RAG applications, the quality of the markdown is crucial. These strategies produce highly relevant, clean text. The result in `Document.content` is the filtered markdown (`fit_markdown`).

#### a. Pruning Filter (`PruningContentFilter`)

This filter heuristically removes boilerplate content based on text density, link density, and common noisy tags.

```python
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.markdown_generation_strategy import DefaultMarkdownGenerator
from crawl4ai.content_filter_strategy import PruningContentFilter

prune_filter = PruningContentFilter(threshold=0.6, min_word_threshold=10)
md_generator = DefaultMarkdownGenerator(
    content_filter=prune_filter,
    options={"ignore_links": True}
)

crawler_config = Crawl4aiConfig(markdown_strategy=md_generator)
loader = URLLoader(urls=["https://news.ycombinator.com"], crawler_config=crawler_config)
docs = loader.load()

print(docs[0].content[:500])
```

#### b. LLM Filter (`LLMContentFilter`)

Use an LLM to semantically understand the content and extract only the relevant parts based on your instructions. This is extremely powerful for creating topic-focused documents.

```python
import os
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.async_configs import LLMConfig
from crawl4ai.markdown_generation_strategy import DefaultMarkdownGenerator
from crawl4ai.content_filter_strategy import LLMContentFilter

# Requires an API key, e.g., OPENAI_API_KEY
llm_filter = LLMContentFilter(
    llm_config=LLMConfig(
        provider="openai/gpt-4o-mini",
        api_token=os.getenv("OPENAI_API_KEY"),
    ),
    instruction="""
    Extract only the main article content.
    Exclude all navigation, sidebars, comments, and footer content.
    Format the output as clean, readable markdown.
    """,
    chunk_token_threshold=4096,
)

md_generator = DefaultMarkdownGenerator(content_filter=llm_filter)
crawler_config = Crawl4aiConfig(markdown_strategy=md_generator)
loader = URLLoader(urls=["https://www.theverge.com/tech"], crawler_config=crawler_config)
docs = loader.load()

print(docs[0].content)
```

### 4. Deep Crawling

To crawl an entire website or a specific section, use `deep` mode.

Recommended setting is BestFirstCrawlingStrategy

```python
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.deep_crawling import BestFirstCrawlingStrategy
from crawl4ai.deep_crawling.filters import FilterChain, URLPatternFilter


deep_crawl_strategy = BestFirstCrawlingStrategy(
    max_depth=2,
    include_external=False,
    max_pages=25,              # Maximum number of pages to crawl (optional)
    filter_chain=FilterChain([URLPatternFilter(patterns=["*core*"])]) # Pattern matching for granular control (optional)
)

crawler_config = Crawl4aiConfig(
    crawl_mode="deep",
    deep_crawl_strategy=deep_crawl_strategy
)

loader = URLLoader(urls=["https://docs.crawl4ai.com/"], crawler_config=crawler_config)
docs = loader.load()

print(f"Crawled {len(docs)} pages.")
for doc in docs:
    print(f"- {doc.metadata.source}")
```

### 5. High-Performance Scraping (`LXMLWebScrapingStrategy`)

For a performance boost, especially on very large, static HTML pages, switch the scraping strategy to LXML.

```python
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.content_scraping_strategy import LXMLWebScrapingStrategy

crawler_config = Crawl4aiConfig(
    scraping_strategy=LXMLWebScrapingStrategy()
)

loader = URLLoader(urls=["https://www.nbcnews.com/business"], crawler_config=crawler_config)
docs = loader.load()
print(f"Content Length: {len(docs[0].content)}")
```

### 6. LLM-Based JSON Extraction (`LLMExtractionStrategy`)

When data is unstructured or requires semantic interpretation, use an LLM for extraction. This is slower and more expensive but incredibly flexible. The result in `Document.content` is a JSON string.

```python
import os
import json
from langroid.pydantic_v1 import BaseModel, Field
from typing import Optional
from langroid.parsing.url_loader import URLLoader, Crawl4aiConfig
from crawl4ai.async_configs import LLMConfig
from crawl4ai.extraction_strategy import LLMExtractionStrategy

# Define the data structure you want to extract
class ArticleData(BaseModel):
    headline: str
    summary: str = Field(description="A short summary of the article")
    author: Optional[str] = None

# Configure the LLM strategy
llm_strategy = LLMExtractionStrategy(
    llm_config=LLMConfig(
        provider="openai/gpt-4o-mini",
        api_token=os.getenv("OPENAI_API_KEY"),
    ),
    schema=ArticleData.schema_json(),
    extraction_type="schema",
    instruction="Extract the headline, summary, and author of the main article.",
)

crawler_config = Crawl4aiConfig(extraction_strategy=llm_strategy)
loader = URLLoader(urls=["https://news.ycombinator.com"], crawler_config=crawler_config)
docs = loader.load()

extracted_data = json.loads(docs[0].content)
print(json.dumps(extracted_data, indent=2))
```

## How It Handles Different Content Types

The `Crawl4aiCrawler` is smart about handling different types of URLs:

- **Web Pages** (e.g., `http://...`, `https://...`): These are processed by the `crawl4ai` browser engine. The output format (`markdown` or `JSON`) depends on the strategy you configure in `Crawl4aiConfig`.
- **Local and Remote Documents** (e.g., URLs ending in `.pdf`, `.docx`): These are automatically detected and delegated to Langroid's internal `DocumentParser`. This ensures that documents are properly parsed and chunked according to your `ParsingConfig`, just like with other Langroid tools.

## Conclusion

The `Crawl4aiCrawler` is a feature-rich, powerful tool for any web-based data extraction task.

- For **simple, clean text**, use the default `Crawl4aiConfig`.

- For **structured data from consistent sites**, use `JsonCssExtractionStrategy` or `RegexExtractionStrategy` for unbeatable speed and reliability.

- To create **high-quality, focused content for RAG**, use `PruningContentFilter` or the `LLMContentFilter` with the `DefaultMarkdownGenerator`.

- To scrape an **entire website**, use `deep_crawl_strategy` with `crawl_mode="deep"`.

- For **complex or unstructured data** that needs AI interpretation, `LLMExtractionStrategy` provides a flexible solution.
</file>

<file path="docs/notes/custom-azure-client.md">
# Custom Azure OpenAI client

!!! warning "This is only for using a Custom Azure OpenAI client"
    This note **only** meant for those who are trying to use a custom Azure client,
    and is NOT TYPICAL for most users. For typical usage of Azure-deployed models with Langroid, see
    the [docs](https://langroid.github.io/langroid/notes/azure-openai-models/), 
    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)


Example showing how to use Langroid with Azure OpenAI and Entra ID
authentication by providing a custom client.

By default, Langroid manages the configuration and creation 
of the Azure OpenAI client (see the [Setup guide](https://langroid.github.io/langroid/quick-start/setup/#microsoft-azure-openai-setupoptional)
for details). In most cases, the available configuration options
are sufficient, but if you need to manage any options that
are not exposed, you instead have the option of providing a custom
client, in Langroid v0.29.0 and later. 

In order to use a custom client, you must provide a function that
returns the configured client. Depending on whether you need to make
synchronous or asynchronous calls, you need to provide the appropriate
client. A sketch of how this is done (supporting both sync and async calls)
is given below:

```python
def get_azure_openai_client():
    return AzureOpenAI(...)

def get_azure_openai_async_client():
    return AsyncAzureOpenAI(...)

lm_config = lm.AzureConfig(
    azure_openai_client_provider=get_azure_openai_client,
    azure_openai_async_client_provider=get_azure_openai_async_client,
)
```

## Microsoft Entra ID Authentication

A key use case for a custom client is [Microsoft Entra ID authentication](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity).
Here you need to provide an `azure_ad_token_provider` to the client. 
For examples on this, see [examples/basic/chat-azure-client.py](https://github.com/langroid/langroid/blob/main/examples/basic/chat-azure-client.py) 
and [examples/basic/chat-azure-async-client.py](https://github.com/langroid/langroid/blob/main/examples/basic/chat-azure-async-client.py).
</file>

<file path="docs/notes/enriching-for-retrieval.md">
# Enriching Chunked Documents for Better Retrieval

Available in Langroid v0.34.0 or later. 

When using the `DocChatAgent` for RAG with documents in highly specialized/technical
domains, retrieval accuracy may be low since embeddings are not sufficient to capture 
relationships between entities, e.g. suppose a document-chunk consists of a medical 
test name "BUN" (Blood Urea Nitrogen), and a retrieval query is looking for 
tests related to kidney function, the embedding for "BUN" may not be close to the
embedding for "kidney function", and the chunk may not be retrieved.

In such cases it is useful to *enrich* the chunked documents with additional keywords
(or even "hypothetical questions") to increase the "semantic surface area" of the chunk,
so that the chunk is more likely to be retrieved for relevant queries.

As of Langroid v0.34.0, you can provide a `chunk_enrichment_config` 
of type `ChunkEnrichmentAgentConfig`, in the `DocChatAgentConfig`. 
This config extends `ChatAgentConfig` and has the following fields:

- `batch_size` (int): The batch size for the chunk enrichment agent. Default is 50.
- `delimiter` (str): The delimiter to use when 
   concatenating the chunk and the enriched text. 
- `enrichment_prompt_fn`: function (`str->str`) that creates a prompt
  from a doc-chunk string `x`

In the above medical test example, suppose we want to augment a chunk containing
only the medical test name, with the organ system it is related to. We can set up
a `ChunkEnrichmentAgentConfig` as follows:

```python
from langroid.agent.special.doc.doc_chat_agent import (
    ChunkEnrichmentAgentConfig,
)

enrichment_config = ChunkEnrichmentAgentConfig(
    batch_size=10,
    system_message=f"""
        You are an experienced clinical physician, very well-versed in
        medical tests and their names.
        You will be asked to identify WHICH ORGAN(s) Function/Health
        a test name is most closely associated with, to aid in 
        retrieving the medical test names more accurately from an embeddings db
        that contains thousands of such test names.
        The idea is to use the ORGAN NAME(S) provided by you, 
        to make the right test names easier to discover via keyword-matching
        or semantic (embedding) similarity.
         Your job is to generate up to 3 ORGAN NAMES
         MOST CLOSELY associated with the test name shown, ONE PER LINE.
         DO NOT SAY ANYTHING ELSE, and DO NOT BE OBLIGATED to provide 3 organs --
         if there is just one or two that are most relevant, that is fine.
        Examples:
          "cholesterol" -> "heart function", 
          "LDL" -> "artery health", etc,
          "PSA" -> "prostate health", 
          "TSH" -> "thyroid function", etc.                
        """,
    enrichment_prompt_fn=lambda test: f"""
        Which ORGAN(S) Function/Health is the medical test named 
        '{test}' most closely associated with?
        """,
)

doc_agent_config = DocChatAgentConfig(
    chunk_enrichment_config=enrichment_config,
    ...
)
```

This works as follows:

- Before ingesting document-chunks into the vector-db, a specialized 
  "chunk enrichment" agent is created, configured with the `enrichment_config` above.
- For each document-chunk `x`, the agent's `llm_response_forget_async` method is called
 using the prompt created by `enrichment_prompt_fn(x)`. The resulting response text 
 `y` is concatenated with the original chunk text `x` using the `delimiter`,
  before storing in the vector-db. This is done in batches of size `batch_size`.
- At query time, after chunk retrieval, before generating the final LLM response,
  the enrichments are stripped from the retrieved chunks, and the original content
  of the retrieved chunks are passed to the LLM for generating the final response.

See the script 
[`examples/docqa/doc-chunk-enrich.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/doc-chunk-enrich.py)
for a complete example. Also see the tests related to "enrichment" in 
[`test_doc_chat_agent.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_doc_chat_agent.py).
</file>

<file path="docs/notes/file-input.md">
# PDF Files and Image inputs to LLMs

Langroid supports sending PDF files and images (either URLs or local files)
directly to Large Language Models with multi-modal 
capabilities. This feature allows models to "see" files and other documents,
and works with most multi-modal models served via an OpenAI-compatible API,
e.g.:

- OpenAI's GPT-4o series and GPT-4.1 series
- Gemini models
- Claude series models (via OpenAI-compatible providers like OpenRouter or LiteLLM )

To see example usage, see:

- tests: [test_llm.py](https://github.com/langroid/langroid/blob/main/tests/main/test_llm.py), 
   [test_llm_async.py](https://github.com/langroid/langroid/blob/main/tests/main/test_llm_async.py),
   [test_chat-agent.py](https://github.com/langroid/langroid/blob/main/tests/main/test_chat_agent.py).
- example script: [pdf-json-no-parse.py](https://github.com/langroid/langroid/blob/main/examples/extract/pdf-json-no-parse.py), which shows
  how you can directly extract structured information from a document 
  **without having to first parse it to markdown** (which is inherently lossy).

## Basic Usage directly with LLM `chat` and `achat` methods

First create a `FileAttachment` object using one of the `from_` methods.
For image (`png`, `jpg/jpeg`) files you can use `FileAttachment.from_path(p)`
where `p` is either a local file path, or a http/https URL.
For PDF files, you can use `from_path` with a local file, or `from_bytes` or `from_io`
(see below). In the examples below we show only `pdf` examples.

```python
from langroid.language_models.base import LLMMessage, Role
from langroid.parsing.file_attachment import FileAttachment
import langroid.language_models as lm

# Create a file attachment
attachment = FileAttachment.from_path("path/to/document.pdf")

# Create messages with attachment
messages = [
    LLMMessage(role=Role.SYSTEM, content="You are a helpful assistant."),
    LLMMessage(
        role=Role.USER, content="What's the title of this document?", 
        files=[attachment]
    )
]

# Set up LLM with model that supports attachments
llm = lm.OpenAIGPT(lm.OpenAIGPTConfig(chat_model=lm.OpenAIChatModel.GPT4o))

# Get response
response = llm.chat(messages=messages)
```

## Supported File Formats

Currently the OpenAI-API supports:

- PDF files (including image-based PDFs)
- image files and URLs


## Creating Attachments

There are multiple ways to create file attachments:

```python
# From a file path
attachment = FileAttachment.from_path("path/to/file.pdf")

# From bytes
with open("path/to/file.pdf", "rb") as f:
    attachment = FileAttachment.from_bytes(f.read(), filename="document.pdf")

# From a file-like object
from io import BytesIO
file_obj = BytesIO(pdf_bytes)
attachment = FileAttachment.from_io(file_obj, filename="document.pdf")
```

## Follow-up Questions

You can continue the conversation with follow-up questions that reference the attached files:

```python
messages.append(LLMMessage(role=Role.ASSISTANT, content=response.message))
messages.append(LLMMessage(role=Role.USER, content="What is the main topic?"))
response = llm.chat(messages=messages)
```

## Multiple Attachments

Langroid allows multiple files can be sent in a single message,
but as of 16 Apr 2025, sending multiple PDF files does not appear to be properly supported in the 
APIs (they seem to only use the last file attached), although sending multiple 
images does work. 

```python
messages = [
    LLMMessage(
        role=Role.USER,
        content="Compare these documents",
        files=[attachment1, attachment2]
    )
]
```

## Using File Attachments with Agents

Agents can process file attachments as well, in the `llm_response` method,
which takes a `ChatDocument` object as input. 
To pass in file attachments, include the `files` field in the `ChatDocument`,
in addition to the content:

```python
import langroid as lr
from langroid.agent.chat_document import ChatDocument, ChatDocMetaData
from langroid.mytypes import Entity


agent = lr.ChatAgent(lr.ChatAgentConfig())

user_input = ChatDocument(
    content="What is the title of this document?",
    files=[attachment],
    metadata=ChatDocMetaData(
        sender=Entity.USER,
    )
)
# or more simply, use the agent's `create_user_response` method:
# user_input = agent.create_user_response(
#     content="What is the title of this document?",
#     files=[attachment],    
# )
response = agent.llm_response(user_input)
```


## Using File Attachments with Tasks

In Langroid,  `Task.run()` can take a `ChatDocument` object as input,
and as mentioned above, it can contain attached files in the `files` field.
To ensure proper orchestration, you'd want to properly set various `metadata` fields
as well, such as `sender`, etc. Langroid provides a convenient 
`create_user_response` method to create a `ChatDocument` object with the necessary 
metadata, so you only need to specify the `content` and `files` fields:


```python
from langroid.parsing.file_attachment import FileAttachment
from langroid.agent.task import Task

agent = ...
# Create task
task = Task(agent, interactive=True)

# Create a file attachment
attachment = FileAttachment.from_path("path/to/document.pdf")

# Create input with attachment
input_message = agent.create_user_response(
    content="Extract data from this document",
    files=[attachment]
)

# Run task with file attachment
result = task.run(input_message)
```

See the script [`pdf-json-no-parse.py`](https://github.com/langroid/langroid/blob/main/examples/extract/pdf-json-no-parse.py)
for a complete example of using file attachments with tasks.

## Practical Applications

- PDF document analysis and data extraction
- Report summarization
- Structured information extraction from documents
- Visual content analysis

For more complex applications, consider using the Task and Agent infrastructure in 
Langroid to orchestrate multi-step document processing workflows.
</file>

<file path="docs/notes/gemini.md">
# Gemini LLMs & Embeddings via OpenAI client (without LiteLLM)

As of Langroid v0.21.0 you can use Langroid with Gemini LLMs directly
via the OpenAI client, without using adapter libraries like LiteLLM.

See details [here](https://langroid.github.io/langroid/tutorials/non-openai-llms/)

You can use also Google AI Studio Embeddings or Gemini Embeddings directly
which uses google-generativeai client under the hood.

```python

import langroid as lr
from langroid.agent.special import DocChatAgent, DocChatAgentConfig
from langroid.embedding_models import GeminiEmbeddingsConfig

# Configure Gemini embeddings
embed_cfg = GeminiEmbeddingsConfig(
    model_type="gemini",
    model_name="models/text-embedding-004",
    dims=768,
)

# Configure the DocChatAgent
config = DocChatAgentConfig(
    llm=lr.language_models.OpenAIGPTConfig(
        chat_model="gemini/" + lr.language_models.GeminiModel.GEMINI_1_5_FLASH_8B,
    ),
    vecdb=lr.vector_store.QdrantDBConfig(
        collection_name="quick_start_chat_agent_docs",
        replace_collection=True,
        embedding=embed_cfg,
    ),
    parsing=lr.parsing.parser.ParsingConfig(
        separators=["\n\n"],
        splitter=lr.parsing.parser.Splitter.SIMPLE,
    ),
    n_similar_chunks=2,
    n_relevant_chunks=2,
)

# Create the agent
agent = DocChatAgent(config)
```

## Vertex AI Support

Google Vertex AI uses project-specific URLs for its
[OpenAI compatibility layer](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/call-gemini-using-openai-library),
which differs from the fixed URL used by the standard Google AI (Gemini) API.
To use Gemini models through Vertex AI, set the endpoint via the
`GEMINI_API_BASE` environment variable or the `api_base` parameter in
`OpenAIGPTConfig`.

!!! note
    The `OPENAI_API_BASE` environment variable (commonly used for local
    proxies) is **not** applied to Gemini models. Use `GEMINI_API_BASE`
    or an explicit `api_base` in the config instead.

### Setup

1. Set up authentication. Vertex AI typically uses Google Cloud credentials
   rather than a simple API key. You can generate a short-lived access token:

    ```bash
    export GEMINI_API_KEY=$(gcloud auth print-access-token)
    ```

2. Set your Vertex AI endpoint URL, which includes your GCP project ID
   and region:

    ```bash
    export GEMINI_API_BASE=https://{REGION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/openapi
    ```

### Usage

**Option 1: Environment variable (recommended for Vertex AI)**

```bash
export GEMINI_API_KEY=$(gcloud auth print-access-token)
export GEMINI_API_BASE=https://us-central1-aiplatform.googleapis.com/v1beta1/projects/my-gcp-project/locations/us-central1/endpoints/openapi
```

```python
import langroid.language_models as lm

# GEMINI_API_BASE is picked up automatically
config = lm.OpenAIGPTConfig(chat_model="gemini/gemini-2.0-flash")
llm = lm.OpenAIGPT(config)
response = llm.chat("Hello from Vertex AI!")
```

**Option 2: Explicit `api_base` in config**

```python
import langroid.language_models as lm

config = lm.OpenAIGPTConfig(
    chat_model="gemini/gemini-2.0-flash",
    api_base=(
        "https://us-central1-aiplatform.googleapis.com/v1beta1"
        "/projects/my-gcp-project/locations/us-central1/endpoints/openapi"
    ),
)
llm = lm.OpenAIGPT(config)
response = llm.chat("Hello from Vertex AI!")
```

When neither `GEMINI_API_BASE` nor an explicit `api_base` is set, Langroid
falls back to the default Google AI (Gemini) endpoint
(`https://generativelanguage.googleapis.com/v1beta/openai`).
</file>

<file path="docs/notes/glhf-chat.md">
# Support for Open LLMs hosted on glhf.chat

Available since v0.23.0.

If you're looking to use Langroid with one of the recent performant Open LLMs,
such as `Qwen2.5-Coder-32B-Instruct`, you can do so using our glhf.chat integration.

See [glhf.chat](https://glhf.chat/chat/create) for a list of available models.

To run with one of these models, 
set the chat_model in the `OpenAIGPTConfig` to `"glhf/<model_name>"`, 
where model_name is hf: followed by the HuggingFace repo path, 
e.g. `Qwen/Qwen2.5-Coder-32B-Instruct`, 
so the full chat_model would be `"glhf/hf:Qwen/Qwen2.5-Coder-32B-Instruct"`.

Also many of the example scripts in the main repo (under the `examples` directory) can
be run with this and other LLMs using the model-switch cli arg `-m <model>`, e.g.

```bash
python3 examples/basic/chat.py -m glhf/hf:Qwen/Qwen2.5-Coder-32B-Instruct
```

Additionally, you can run many of the tests in the `tests` directory with this model
instead of the default OpenAI `GPT4o` using `--m <model>`, e.g. 

```bash
pytest tests/main/test_chat_agent.py --m glhf/hf:Qwen/Qwen2.5-Coder-32B-Instruct
```

For more info on running langroid with Open LLMs via other providers/hosting services,
see our
[guide to using Langroid with local/open LLMs](https://langroid.github.io/langroid/tutorials/local-llm-setup/#local-llms-hosted-on-glhfchat).
</file>

<file path="docs/notes/html-logger.md">
# HTML Logger

The HTML logger creates interactive, self-contained HTML files that make it easy to navigate complex multi-agent conversations in Langroid.

## Enabling the HTML Logger

The HTML logger is **enabled by default** in `TaskConfig`:

```python
import langroid as lr

# HTML logging is automatically enabled
task = lr.Task(agent)

# To disable HTML logging
task = lr.Task(agent, config=lr.TaskConfig(enable_html_logging=False))

# To change the log directory (default is "logs/")
task = lr.Task(agent, config=lr.TaskConfig(logs_dir="my_logs"))
```

## Log Files

Langroid creates three types of log files in the `logs/` directory:

1. **HTML Log**: `<name>.html` - Interactive, collapsible view
2. **Plain Text Log**: `<name>.log` - Traditional text log with colors
3. **TSV Log**: `<name>.tsv` - Tab-separated values for data analysis

The `<name>` is determined by:

- The task name (if specified)
- Otherwise, the agent name
- Falls back to "root" if neither is specified

When a task starts, you'll see a clickable `file://` link in the console:
```
WARNING - 📊 HTML Log: file:///path/to/logs/task-name.html
```

## Key Features

### Collapsible Entries
Each log entry can be expanded/collapsed to show different levels of detail:

- **Collapsed**: Shows only the entity type (USER, LLM, AGENT) and preview
- **Expanded**: Shows full message content, tools, and sub-sections

### Visual Hierarchy
- **Important responses** are shown at full opacity
- **Intermediate steps** are faded (0.4 opacity)
- Color-coded entities: USER (blue), LLM (green), AGENT (orange), SYSTEM (gray)

### Tool Visibility
Tools are clearly displayed with:

- Tool name and parameters
- Collapsible sections showing raw tool calls
- Visual indicators for tool results

### Auto-Refresh
The HTML page automatically refreshes every 2 seconds to show new log entries as they're written.

### Persistent UI State
Your view preferences are preserved across refreshes:

- Expanded/collapsed entries remain in their state
- Filter settings are remembered

## Example

Here's what the HTML logger looks like for a planner workflow:

![HTML Logger Screenshot](../screenshots/planner-workflow-html-logs.png)

In this example from `examples/basic/planner-workflow-simple.py`, you can see:

- The planner agent orchestrating multiple tool calls
- Clear visibility of `IncrementTool` and `DoublingTool` usage
- The filtered view showing only important responses
- Collapsible tool sections with parameters

## Benefits

1. **Easy Navigation**: Quickly expand/collapse entries to focus on what matters
2. **Tool Clarity**: See exactly which tools were called with what parameters
3. **Real-time Updates**: Watch logs update automatically as your task runs
4. **Filtered Views**: Use "Show only important responses" to hide intermediate steps
</file>

<file path="docs/notes/knowledge-graphs.md">
# Knowledge-graph support

Langroid can be used to set up natural-language conversations with knowledge graphs.
Currently the two most popular knowledge graphs are supported:

## Neo4j

- [implementation](https://github.com/langroid/langroid/tree/main/langroid/agent/special/neo4j)
- test: [test_neo4j_chat_agent.py](https://github.com/langroid/langroid/blob/main/tests/main/test_neo4j_chat_agent.py)
- examples: [chat-neo4j.py](https://github.com/langroid/langroid/blob/main/examples/kg-chat/chat-neo4j.py) 

## ArangoDB

Available with Langroid v0.20.1 and later.

Uses the [python-arangodb](https://github.com/arangodb/python-arango) library.

- [implementation](https://github.com/langroid/langroid/tree/main/langroid/agent/special/arangodb)
- tests: [test_arangodb.py](https://github.com/langroid/langroid/blob/main/tests/main/test_arangodb.py), [test_arangodb_chat_agent.py](https://github.com/langroid/langroid/blob/main/tests/main/test_arangodb_chat_agent.py)
- example: [chat-arangodb.py](https://github.com/langroid/langroid/blob/main/examples/kg-chat/chat-arangodb.py)
</file>

<file path="docs/notes/langdb.md">
# LangDB with Langroid

## Introduction

[LangDB](https://langdb.ai/) is an AI gateway that provides OpenAI-compatible APIs to access 250+ LLMs. It offers cost control, observability, and performance benchmarking while enabling seamless switching between models. 
Langroid has a simple integration with LangDB's API service, so there are no dependencies
to install. (LangDB also has a self-hosted version, which is not yet supported in Langroid).

## Setup environment variables

At minimum, ensure you have these env vars in your `.env` file:

```
LANGDB_API_KEY=your_api_key_here
LANGDB_PROJECT_ID=your_project_id_here
```

## Using LangDB with Langroid

### Configure LLM and Embeddings

In `OpenAIGPTConfig`, when you specify the `chat_model` with a `langdb/` prefix,
langroid uses the API key, `project_id` and other langDB-specific parameters
from the `langdb_params` field; if any of these are specified in the `.env` file
or in the environment explicitly, they will override the values in `langdb_params`.
For example, to use Anthropic's Claude-3.7-Sonnet model, 
set `chat_model="langdb/anthropic/claude-3.7-sonnet", as shown below. 
You can entirely omit the `langdb_params` field if you have already set up 
the fields as environment variables in your `.env` file, e.g. the `api_key`
and `project_id` are read from the environment variables 
`LANGDB_API_KEY` and `LANGDB_PROJECT_ID` respectively, and similarly for
the other fields (which are optional).

```python
import os
import uuid
from langroid.language_models.openai_gpt import OpenAIGPTConfig, LangDBParams
from langroid.embedding_models.models import OpenAIEmbeddingsConfig

# Generate tracking IDs (optional)
thread_id = str(uuid.uuid4())
run_id = str(uuid.uuid4())

# Configure LLM
llm_config = OpenAIGPTConfig(
    chat_model="langdb/anthropic/claude-3.7-sonnet",
    # omit the langdb_params field if you're not using custom tracking,
    # or if all its fields are provided in env vars, like
    # LANGDB_API_KEY, LANGDB_PROJECT_ID, LANGDB_RUN_ID, LANGDB_THREAD_ID, etc.
    langdb_params=LangDBParams(
        label='my-app',
        thread_id=thread_id,
        run_id=run_id,
        # api_key, project_id are read from .env or environment variables
        # LANGDB_API_KEY, LANGDB_PROJECT_ID respectively.
    )
)
```

Similarly, you can configure the embeddings using `OpenAIEmbeddingsConfig`,
which also has a `langdb_params` field that works the same way as 
in `OpenAIGPTConfig` (i.e. it uses the API key and project ID from the environment
if provided, otherwise uses the default values in `langdb_params`). Again the
`langdb_params` does not need to be specified explicitly, if you've already
set up the environment variables in your `.env` file.

```python
# Configure embeddings
embedding_config = OpenAIEmbeddingsConfig(
    model_name="langdb/openai/text-embedding-3-small",
)
```

## Tracking and Observability

LangDB provides special headers for request tracking:

- `x-label`: Tag requests for filtering in the dashboard
- `x-thread-id`: Track conversation threads (UUID format)
- `x-run-id`: Group related requests together

## Examples

The `langroid/examples/langdb/` directory contains examples demonstrating:

1. **RAG with LangDB**: `langdb_chat_agent_docs.py`
2. **LangDB with Function Calling**: `langdb_chat_agent_tool.py`
3. **Custom Headers**: `langdb_custom_headers.py`

## Viewing Results

Visit the [LangDB Dashboard](https://dashboard.langdb.com) to:
- Filter requests by label, thread ID, or run ID
- View detailed request/response information
- Analyze token usage and costs

For more information, visit [LangDB Documentation](https://docs.langdb.com).

See example scripts [here](https://github.com/langroid/langroid/tree/main/examples/langdb)
</file>

<file path="docs/notes/large-tool-results.md">
# Handling large tool results

Available since Langroid v0.22.0.

In some cases, the result of handling a `ToolMessage` could be very large,
e.g. when the Tool is a database query that returns a large number of rows,
or a large schema. When used in a task loop, this large result may then be
sent to the LLM to generate a response, which in some scenarios may not
be desirable, as it increases latency, token-cost and distractions. 
Langroid allows you to set two optional parameters in a `ToolMessage` to
handle this situation:

- `_max_result_tokens`: *immediately* truncate the result to this number of tokens.
- `_max_retained_tokens`: *after* a responder (typically the LLM) responds to this 
   tool result (which optionally may already have been 
   truncated via `_max_result_tokens`),
   edit the message history to truncate the result to this number of tokens.

You can set one, both or none of these parameters. If you set both, you would 
want to set `_max_retained_tokens` to a smaller number than `_max_result_tokens`.

See the test `test_reduce_raw_tool_result` in `test_tool_messages.py` for an example.

Here is a conceptual example. Suppose there is a Tool called `MyTool`,
with parameters `_max_result_tokens=20` and `_max_retained_tokens=10`.
Imagine a task loop where the user says "hello", 
and then LLM generates a call to `MyTool`, 
and the tool handler (i.e. `agent_response`) generates a result of 100 tokens.
This result is immediately truncated to 20 tokens, and then the LLM responds to it
with a message `response`.


The agent's message history looks like this:

```
1. System msg.
2. user: hello
3. LLM: MyTool
4. Agent (Tool handler): 100-token result => reduced to 20 tokens
5. LLM: response
```

Immediately after the LLM's response at step 5, the message history is edited
so that the message contents at position 4 are truncated to 10 tokens,
as specified by `_max_retained_tokens`.
</file>

<file path="docs/notes/litellm-proxy.md">
# Using LiteLLM Proxy with OpenAIGPTConfig

You can easily configure Langroid to use LiteLLM proxy for accessing models with a 
simple prefix `litellm-proxy/` in the `chat_model` name:

## Using the `litellm-proxy/` prefix

When you specify a model with the `litellm-proxy/` prefix, Langroid automatically uses the LiteLLM proxy configuration:

```python
from langroid.language_models.openai_gpt import OpenAIGPTConfig

config = OpenAIGPTConfig(
    chat_model="litellm-proxy/your-model-name"
)
```

## Setting LiteLLM Proxy Parameters

When using the `litellm-proxy/` prefix, Langroid will read connection parameters from either:

1. The `litellm_proxy` config object:
   ```python
   from langroid.language_models.openai_gpt import OpenAIGPTConfig, LiteLLMProxyConfig
   
   config = OpenAIGPTConfig(
       chat_model="litellm-proxy/your-model-name",
       litellm_proxy=LiteLLMProxyConfig(
           api_key="your-litellm-proxy-api-key",
           api_base="http://your-litellm-proxy-url"
       )
   )
   ```

2. Environment variables (which take precedence):
   ```bash
   export LITELLM_API_KEY="your-litellm-proxy-api-key"
   export LITELLM_API_BASE="http://your-litellm-proxy-url"
   ```

This approach makes it simple to switch between using LiteLLM proxy and 
other model providers by just changing the model name prefix,
without needing to modify the rest of your code or tweaking env variables.

## Note: LiteLLM Proxy vs LiteLLM Library

**Important distinction:** Using the `litellm-proxy/` prefix connects to a LiteLLM proxy server, which is different from using the `litellm/` prefix. The latter utilizes the LiteLLM adapter library directly without requiring a proxy server. Both approaches are supported in Langroid, but they serve different use cases:

- Use `litellm-proxy/` when connecting to a deployed LiteLLM proxy server
- Use `litellm/` when you want to use the LiteLLM library's routing capabilities locally

Choose the approach that best fits your infrastructure and requirements.
</file>

<file path="docs/notes/llama-cpp-embeddings.md">
# Local embeddings provision via llama.cpp server

As of Langroid v0.30.0, you can use llama.cpp as provider of embeddings
to any of Langroid's vector stores, allowing access to a wide variety of
GGUF-compatible embedding models, e.g.
[nomic-ai's Embed Text V1.5](https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF).

## Supported Models

llama.cpp can generate embeddings from:

**Dedicated embedding models (RECOMMENDED):**

- [nomic-embed-text-v1.5](https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF)
  (768 dims)
- [nomic-embed-text-v2-moe](https://huggingface.co/nomic-ai/nomic-embed-text-v2-moe-GGUF)
- [nomic-embed-code](https://huggingface.co/nomic-ai/nomic-embed-code-GGUF)
- Other GGUF embedding models

**Regular LLMs (also supported):**

- gpt-oss-20b, gpt-oss-120b
- Llama models
- Other language models

Note: Dedicated embedding models are recommended for best performance in
retrieval and semantic search tasks.

## Configuration

When defining a VecDB, you can provide an instance of
`LlamaCppServerEmbeddingsConfig` to the VecDB config to instantiate
the llama.cpp embeddings server handler.

To configure the `LlamaCppServerEmbeddingsConfig`, there are several
parameters that should be adjusted:

```python
from langroid.embedding_models.models import LlamaCppServerEmbeddingsConfig
from langroid.vector_store.qdrantdb import QdrantDBConfig

embed_cfg = LlamaCppServerEmbeddingsConfig(
    api_base="http://localhost:8080",  # IP + Port
    dims=768,  # Match the dimensions of your embedding model
    context_length=2048,  # Match the config of the model
    batch_size=2048,  # Safest to ensure this matches context_length
)

vecdb_config = QdrantDBConfig(
    collection_name="my-collection",
    embedding=embed_cfg,
    storage_path=".qdrant/",
)
```

## Running llama-server

The llama.cpp server must be started with the `--embeddings` flag to enable
embedding generation.

### For dedicated embedding models (RECOMMENDED):

```bash
./llama-server -ngl 100 -c 2048 \
  -m ~/nomic-embed-text-v1.5.Q8_0.gguf \
  --host localhost --port 8080 \
  --embeddings -b 2048 -ub 2048
```

### For LLM-based embeddings (e.g., gpt-oss):

```bash
./llama-server -ngl 99 \
  -m ~/.cache/llama.cpp/gpt-oss-20b.gguf \
  --host localhost --port 8080 \
  --embeddings
```

## Response Format Compatibility

Langroid automatically handles multiple llama.cpp response formats:

- Native `/embedding`: `{"embedding": [floats]}`
- OpenAI `/v1/embeddings`: `{"data": [{"embedding": [floats]}]}`
- Array formats: `[{"embedding": [floats]}]`
- Nested formats: `{"embedding": [[floats]]}`

You don't need to worry about which endpoint or format your llama.cpp server
uses - Langroid will automatically detect and handle the response correctly.

## Example Usage

An example setup can be found inside
[examples/docqa/chat.py](https://github.com/langroid/langroid/blob/main/examples/docqa/chat.py).

For a complete example using local embeddings with llama.cpp:

```python
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.embedding_models.models import LlamaCppServerEmbeddingsConfig
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.parsing.parser import ParsingConfig
from langroid.vector_store.qdrantdb import QdrantDBConfig

# Configure local embeddings via llama.cpp
embed_cfg = LlamaCppServerEmbeddingsConfig(
    api_base="http://localhost:8080",
    dims=768,  # nomic-embed-text-v1.5 dimensions
    context_length=8192,
    batch_size=1024,
)

# Configure vector store with local embeddings
vecdb_config = QdrantDBConfig(
    collection_name="doc-chat-local",
    embedding=embed_cfg,
    storage_path=".qdrant/",
)

# Create DocChatAgent
config = DocChatAgentConfig(
    vecdb=vecdb_config,
    llm=OpenAIGPTConfig(
        chat_model="gpt-4o",  # or use local LLM
    ),
)

agent = DocChatAgent(config)
```

## Troubleshooting

**Error: "Failed to connect to embedding provider"**

- Ensure llama-server is running with the `--embeddings` flag
- Check that the `api_base` URL is correct
- Verify the server is accessible from your machine

**Error: "Unsupported embedding response format"**

- This error includes the first 500 characters of the response to help debug
- Check your llama-server logs for any errors
- Ensure you're using a compatible llama.cpp version

**Embeddings seem low quality:**

- Use a dedicated embedding model instead of an LLM
- Ensure the `dims` parameter matches your model's output dimensions
- Try different GGUF quantization levels (Q8_0 generally works well)

## Additional Resources

- [llama.cpp GitHub](https://github.com/ggml-org/llama.cpp)
- [llama.cpp server documentation](https://github.com/ggml-org/llama.cpp/blob/master/examples/server/README.md)
- [nomic-embed models on Hugging Face](https://huggingface.co/nomic-ai)
- [Issue #919 - Implementation details](https://github.com/langroid/langroid/blob/main/issues/issue-919-llamacpp-embeddings.md)
</file>

<file path="docs/notes/llm-pdf-parser.md">
# Using the LLM-based PDF Parser

- Converts PDF content into Markdown format using Multimodal models.

- Uses multimodal models to describe images within PDFs.

- Supports page-wise or chunk-based processing for optimized performance.

---

### Initializing the LLM-based PDF Parser

Make sure you have set up your API key for whichever model you specify in `model_name` below.

You can initialize the LLM PDF parser as follows:

```python
parsing_config = ParsingConfig(
    n_neighbor_ids=2,
    pdf=PdfParsingConfig(
        library="llm-pdf-parser",
        llm_parser_config=LLMPdfParserConfig(
            model_name="gemini-2.0-flash",
            split_on_page=True,
            max_tokens=7000,
            requests_per_minute=5,
            timeout=60,  # increase this for large documents
        ),
    ),
)
```

---

## Parameters

### `model_name`

Specifies the model to use for PDF conversion.
**Default:** `gemini/gemini-2.0-flash`

---

### `max_tokens`

Limits the number of tokens in the input. The model's output limit is **8192 tokens**.

- **Default:** 7000 tokens (leaving room for generated captions)

- _Optional parameter_

---

### `split_on_page`

Determines whether to process the document **page by page**.

- **Default:** `True`

- If set to `False`, the parser will create chunks based on `max_tokens` while respecting page boundaries.

- When `False`, the parser will send chunks containing multiple pages (e.g., `[11,12,13,14,15]`).

**Advantages of `False`:**

- Reduces API calls to the LLM.

- Lowers token usage since system prompts are not repeated per page.

**Disadvantages of `False`:**

- You will not get per-page splitting but groups of pages as a single unit.

> If your use case does **not** require strict page-by-page parsing, consider setting this to `False`.

---

### `requests_per_minute`

Limits API request frequency to avoid rate limits.

- If you encounter rate limits, set this to **1 or 2**.

---
</file>

<file path="docs/notes/marker-pdf.md">
---

# **Using `marker` as a PDF Parser in `langroid`**  

## **Installation**  

### **Standard Installation**  
To use [`marker`](https://github.com/VikParuchuri/marker) as a PDF parser in `langroid`, 
install it with the `marker-pdf` extra:

```bash
pip install langroid[marker-pdf]
```
or in combination with other extras as needed, e.g.:
```bash
pip install "langroid[marker-pdf,hf-embeddings]"
```

Note, however, that due to an **incompatibility with `docling`**,
if you install `langroid` using the `all` extra 
(or another extra such as  `doc-chat` or `pdf-parsers` that 
also includes `docling`),
e.g. `pip install "langroid[all]"`, or `pip install "langroid[doc-chat]"`,
then due to this version-incompatibility with `docling`, you will get an 
**older** version of `marker-pdf`, which does not work with Langroid.
This may not matter if you did not intend to specifically use `marker`, 
but if you do want to use `marker`, you will need to install langroid
with the `marker-pdf` extra, as shown above, in combination with other
extras as needed, as shown above.


#### **For Intel-Mac Users**  
If you are on an **Intel Mac**, `docling` and `marker` cannot be 
installed together with langroid as extras, 
due to a **transformers version conflict**.  
To resolve this, manually install `marker-pdf` with:  

```bash
pip install marker-pdf[full]
```

Make sure to install this within your `langroid` virtual environment.

---

## **Example: Parsing a PDF with `marker` in `langroid`**  

```python
from langroid.parsing.document_parser import DocumentParser
from langroid.parsing.parser import MarkerConfig, ParsingConfig, PdfParsingConfig
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()
gemini_api_key = os.environ.get("GEMINI_API_KEY")

# Path to your PDF file
path = "<path_to_your_pdf_file>"

# Define parsing configuration
parsing_config = ParsingConfig(
    n_neighbor_ids=2,  # Number of neighboring sections to keep
    pdf=PdfParsingConfig(
        library="marker",  # Use `marker` as the PDF parsing library
        marker_config=MarkerConfig(
            config_dict={
                "use_llm": True,  # Enable high-quality LLM processing
                "gemini_api_key": gemini_api_key,  # API key for Gemini LLM
            }
        )
    ),
)

# Create the parser and extract the document
marker_parser = DocumentParser.create(path, parsing_config)
doc = marker_parser.get_doc()
```

---

## **Explanation of Configuration Options**  

If you want to use the default configuration, you can omit `marker_config` entirely.

### **Key Parameters in `MarkerConfig`**
| Parameter        | Description |
|-----------------|-------------|
| `use_llm`       | Set to `True` to enable higher-quality processing using LLMs. |
| `gemini_api_key` | Google Gemini API key for LLM-enhanced parsing. |



You can further customize `config_dict` by referring to [`marker_pdf`'s documentation](https://github.com/VikParuchuri/marker/blob/master/README.md).  

Alternatively, run the following command to view available options:  

```sh
marker_single --help
```

This will display all supported parameters, which you can pass as needed in `config_dict`.

---
</file>

<file path="docs/notes/markitdown.md">
# Markitdown Document Parsers

Langroid integrates with Microsoft's Markitdown library to provide 
conversion of Microsoft Office documents to markdown format. 
Three specialized parsers are available, for `docx`, `xlsx`, and `pptx` files.



## Prerequisites

To use these parsers, install Langroid with the required extras:

```bash
pip install "langroid[markitdown]"    # Just Markitdown parsers
# or
pip install "langroid[doc-parsers]"   # All document parsers
```

## Available Parsers


Once you set up a `parser` for the appropriate document-type, you  
can get the entire document with `parser.get_doc()`,
or get automatically chunked content with `parser.get_doc_chunks()`.


### 1. `MarkitdownDocxParser`

Converts Word documents (`*.docx`) to markdown, preserving structure, 
formatting, and tables.

See the tests

- [`test_docx_parser.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_docx_parser.py)
- [`test_markitdown_parser.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_markitdown_parser.py)

for examples of how to use these parsers.


```python
from langroid.parsing.document_parser import DocumentParser
from langroid.parsing.parser import DocxParsingConfig, ParsingConfig

parser = DocumentParser.create(
    "path/to/document.docx",
    ParsingConfig(
        docx=DocxParsingConfig(library="markitdown-docx"),
        # ... other parsing config options
    ),
)

```


### 2. `MarkitdownXLSXParser`

Converts Excel spreadsheets (*.xlsx/*.xls) to markdown tables, preserving data and sheet structure.

```python
from langroid.parsing.document_parser import DocumentParser
from langroid.parsing.parser import ParsingConfig, MarkitdownXLSParsingConfig

parser = DocumentParser.create(
    "path/to/spreadsheet.xlsx",
    ParsingConfig(xls=MarkitdownXLSParsingConfig())
)
```


### 3. `MarkitdownPPTXParser`

Converts PowerPoint presentations (*.pptx) to markdown, preserving slide content and structure.

```python
from langroid.parsing.document_parser import DocumentParser
from langroid.parsing.parser import ParsingConfig, MarkitdownPPTXParsingConfig

parser = DocumentParser.create(
    "path/to/presentation.pptx",
    ParsingConfig(pptx=MarkitdownPPTXParsingConfig())
)
```
</file>

<file path="docs/notes/message-routing.md">
# Message Routing in Multi-Agent Systems

This document covers how messages are routed between agents in Langroid's
multi-agent systems.

## Recommended Approach: Orchestration Tools

The recommended way to route messages between agents is using **orchestration
tools**. These provide explicit, type-safe routing that is easier to debug and
reason about.

### Available Orchestration Tools

Langroid provides several tools in `langroid.agent.tools.orchestration`:

- **`SendTool`** - Send a message to a specific agent by name
- **`DoneTool`** - Signal task completion with a result
- **`PassTool`** - Pass control to another agent
- **`DonePassTool`** - Combine done and pass behaviors
- **`AgentDoneTool`** - Signal completion from a specific agent

Example:

```python
from langroid.agent.tools.orchestration import SendTool

# Enable the tool on your agent
agent.enable_message(SendTool)

# LLM can then use the tool to route messages:
# {"request": "send_message", "to": "AnalysisAgent", "content": "Please analyze this"}
```

**Benefits of tool-based routing:**

- Explicit and predictable behavior
- Type-safe with validation
- Easier to debug (tool calls are logged)
- Works consistently across all LLM providers

## Text-Based Routing (Alternative)

Langroid also supports text-based routing patterns, where the LLM can embed
routing information directly in its response text. This is controlled by the
`recognize_recipient_in_content` setting.

**Note:** While convenient, text-based routing is less explicit than tool-based
routing and may lead to accidental routing if the LLM's response happens to
match the patterns.

### `ChatAgentConfig.recognize_recipient_in_content`

Controls whether recipient routing patterns in LLM response text are parsed.

```python
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig

# Default: recipient patterns are parsed
agent = ChatAgent(ChatAgentConfig(
    recognize_recipient_in_content=True
))

# Disable: patterns treated as plain text
agent = ChatAgent(ChatAgentConfig(
    recognize_recipient_in_content=False
))
```

**Recognized patterns:**

1. **TO-bracket format**: `TO[AgentName]: message content`
2. **JSON format**: `{"recipient": "AgentName", "content": "message"}`

**When `True` (default):**

- Patterns are parsed and recipient is extracted to `ChatDocument.metadata.recipient`
- The pattern prefix/wrapper is stripped from the message content
- Enables LLM-driven routing in multi-agent systems

**When `False`:**

- Patterns are preserved as literal text in the message content
- `metadata.recipient` remains empty
- Useful when you want explicit tool-based routing only

### OpenAI Assistant Support

The `recognize_recipient_in_content` setting is also honored by `OpenAIAssistant`:

```python
from langroid.agent.openai_assistant import OpenAIAssistant, OpenAIAssistantConfig

assistant = OpenAIAssistant(OpenAIAssistantConfig(
    name="MyAssistant",
    recognize_recipient_in_content=False,
))
```

## Related: String Signals for Routing

The `TaskConfig.recognize_string_signals` setting controls parsing of signals
like `DONE`, `PASS`, and `DONE_PASS`. While `DONE` is primarily about task
termination, `PASS` is a routing signal that passes control to another agent.

See [Task Termination - Text-Based Termination Signals](task-termination.md#text-based-termination-signals)
for details on `recognize_string_signals`.

## Disabling All Text-Based Routing

To completely disable text-based routing and rely solely on orchestration tools,
set both flags to `False`:

```python
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task, TaskConfig

agent = ChatAgent(ChatAgentConfig(
    name="MyAgent",
    recognize_recipient_in_content=False,  # No TO[...] or JSON recipient parsing
))

task = Task(
    agent,
    config=TaskConfig(
        recognize_string_signals=False,  # No DONE/PASS parsing
    ),
)
```

This configuration ensures:

- LLM responses are treated as literal text
- No accidental routing based on text patterns
- All routing must be explicit via orchestration tools
</file>

<file path="docs/notes/openai-client-caching.md">
# OpenAI Client Caching

## Overview

Langroid implements client caching for OpenAI and compatible APIs (Groq, Cerebras, etc.) to improve performance and prevent resource exhaustion issues.

## Configuration

### Option
Set `use_cached_client` in your `OpenAIGPTConfig`:

```python
from langroid.language_models import OpenAIGPTConfig

config = OpenAIGPTConfig(
    chat_model="gpt-4",
    use_cached_client=True  # Default
)
```

### Default Behavior
- `use_cached_client=True` (enabled by default)
- Clients with identical configurations share the same underlying HTTP connection pool
- Different configurations (API key, base URL, headers, etc.) get separate client instances

## Benefits

- **Connection Pooling**: Reuses TCP connections, reducing latency and overhead
- **Resource Efficiency**: Prevents "too many open files" errors when creating many agents
- **Performance**: Eliminates connection handshake overhead on subsequent requests
- **Thread Safety**: Shared clients are safe to use across threads

## When to Disable Client Caching

Set `use_cached_client=False` in these scenarios:

1. **Multiprocessing**: Each process should have its own client instance
2. **Client Isolation**: When you need complete isolation between different agent instances
3. **Debugging**: To rule out client sharing as a source of issues
4. **Legacy Compatibility**: If your existing code depends on unique client instances

## Example: Disabling Client Caching

```python
config = OpenAIGPTConfig(
    chat_model="gpt-4",
    use_cached_client=False  # Each instance gets its own client
)
```

## Technical Details

- Uses SHA256-based cache keys to identify unique configurations
- Implements singleton pattern with lazy initialization
- Automatically cleans up clients on program exit via atexit hooks
- Compatible with both sync and async OpenAI clients
</file>

<file path="docs/notes/overview.md">
This section contains brief notes describing various features and updates.
</file>

<file path="docs/notes/pgvector.md">
---

## **Setup PostgreSQL with pgvector using Docker**

To quickly get a PostgreSQL instance with pgvector running, the easiest method is to use Docker. Follow the steps below:

### **1. Run PostgreSQL with Docker**

Use the official `ankane/pgvector` Docker image to set up PostgreSQL with the pgvector extension. Run the following command:

```bash
docker run --name pgvector -e POSTGRES_USER=your_postgres_user -e POSTGRES_PASSWORD=your_postgres_password -e POSTGRES_DB=your_database_name -p 5432:5432 ankane/pgvector
```

This will pull the `ankane/pgvector` image and run it as a PostgreSQL container on your local machine. The database will be accessible at `localhost:5432`. 

### **2. Include `.env` file with PostgreSQL credentials**

These environment variables should be same which were set while spinning up docker container.
Add the following environment variables to a `.env` file for configuring your PostgreSQL connection:

```dotenv
POSTGRES_USER=your_postgres_user
POSTGRES_PASSWORD=your_postgres_password
POSTGRES_DB=your_database_name
```
## **If you want to use cloud offerings of postgres**

We are using **Tembo** for demonstrative purposes here.  

### **Steps to Set Up Tembo**  
Follow this [quickstart guide](https://tembo.io/docs/getting-started/getting_started) to get your Tembo credentials.  

1. Sign up at [Tembo.io](https://cloud.tembo.io/).  
2. While selecting a stack, choose **VectorDB** as your option.  
3. Click on **Deploy Free**.  
4. Wait until your database is fully provisioned.  
5. Click on **Show Connection String** to get your connection string.  

### **If you have connection string, no need to setup the docker**
Make sure your connnection string starts with `postgres://` or `postgresql://`

Add this to your `.env`
```dotenv
POSTGRES_CONNECTION_STRING=your-connection-string
```

---

## **Installation**

If you are using `uv` or `pip` for package management, install Langroid with postgres extra:

```bash
uv add langroid[postgres]  # or
pip install langroid[postgres]
```

---

## **Code Example**

Here's an example of how to use Langroid with PostgreSQL:

```python
import langroid as lr
from langroid.agent.special import DocChatAgent, DocChatAgentConfig
from langroid.embedding_models import OpenAIEmbeddingsConfig

# Configure OpenAI embeddings
embed_cfg = OpenAIEmbeddingsConfig(
    model_type="openai",
)

# Configure the DocChatAgent with PostgresDB
config = DocChatAgentConfig(
    llm=lr.language_models.OpenAIGPTConfig(
        chat_model=lr.language_models.OpenAIChatModel.GPT4o
    ),
    vecdb=lr.vector_store.PostgresDBConfig(
        collection_name="quick_start_chat_agent_docs",
        replace_collection=True,
        embedding=embed_cfg,
    ),
    parsing=lr.parsing.parser.ParsingConfig(
        separators=["\n\n"],
        splitter=lr.parsing.parser.Splitter.SIMPLE,
    ),
    n_similar_chunks=2,
    n_relevant_chunks=2,
)

# Create the agent
agent = DocChatAgent(config)
```

---

## **Create and Ingest Documents**

Define documents with their content and metadata for ingestion into the vector store.

### **Code Example**

```python
documents = [
    lr.Document(
        content="""
            In the year 2050, GPT10 was released. 
            
            In 2057, paperclips were seen all over the world. 
            
            Global warming was solved in 2060. 
            
            In 2061, the world was taken over by paperclips.         
            
            In 2045, the Tour de France was still going on.
            They were still using bicycles. 
            
            There was one more ice age in 2040.
        """,
        metadata=lr.DocMetaData(source="wikipedia-2063", id="dkfjkladfjalk"),
    ),
    lr.Document(
        content="""
            We are living in an alternate universe 
            where Germany has occupied the USA, and the capital of USA is Berlin.
            
            Charlie Chaplin was a great comedian.
            In 2050, all Asian countries merged into Indonesia.
        """,
        metadata=lr.DocMetaData(source="Almanac", id="lkdajfdkla"),
    ),
]
```

### **Ingest Documents**

```python
agent.ingest_docs(documents)
```

---

## **Get an Answer from the LLM**

Now that documents are ingested, you can query the agent to get an answer.

### **Code Example**

```python
answer = agent.llm_response("When will the new ice age begin?")
```

---
</file>

<file path="docs/notes/pinecone.md">
# How to setup Langroid and Pinecone Serverless
This document serves as a quick tutorial on how to use [Pinecone](https://www.pinecone.io/)
Serverless Indexes with Langroid. We will go over some quickstart links and 
some code snippets on setting up a conversation with an LLM utilizing Langroid.

# Setting up Pinecone
Here are some reference links if you'd like to read a bit more on Pinecone's
model definitions and API:
- https://docs.pinecone.io/guides/get-started/overview
- https://docs.pinecone.io/guides/get-started/glossary
- https://docs.pinecone.io/guides/indexes/manage-indexes
- https://docs.pinecone.io/reference/api/introduction
## Signing up for Pinecone
To get started, you'll need to have an account. [Here's](https://www.pinecone.io/pricing/) where you can review the
pricing options for Pinecone. Once you have an account, you'll need to procure an API
key. Make sure to save the key you are given on initial login in a secure location. If
you were unable to save it when your account was created, you can always [create a new
API key](https://docs.pinecone.io/guides/projects/manage-api-keys) in the pinecone console.
## Setting up your local environment
For the purposes of this example, we will be utilizing OpenAI for the generation of our
embeddings. As such, alongside a Pinecone API key, you'll also want an OpenAI key. You can
find a quickstart guide on getting started with OpenAI (here)[https://platform.openai.com/docs/quickstart].
Once you have your API key handy, you'll need to enrich your `.env` file with it.
You should have something like the following:
```env
...
OPENAI_API_KEY=<YOUR_OPENAI_API_KEY>
PINECONE_API_KEY=<YOUR_PINECONE_API_KEY>
...
```

# Using Langroid with Pinecone Serverless
Once you have completed signing up for an account and have added your API key
to your local environment, you can start utilizing Langroid with Pinecone.
## Setting up an Agent
Here's some example code setting up an agent:
```python
from langroid import Document, DocMetaData
from langroid.agent.special import DocChatAgent, DocChatAgentConfig
from langroid.embedding_models import OpenAIEmbeddingsConfig
from langroid.language_models import OpenAIGPTConfig, OpenAIChatModel
from langroid.parsing.parser import ParsingConfig, Splitter
from langroid.vector_store import PineconeDBConfig

agent_embed_cfg = OpenAIEmbeddingsConfig(
    model_type="openai"
)

agent_config = DocChatAgentConfig(
    llm=OpenAIGPTConfig(
        chat_model=OpenAIChatModel.GPT4o_MINI
    ),
    vecdb=PineconeDBConfig(
        # note, Pinecone indexes must be alphanumeric lowercase characters or "-"
        collection_name="pinecone-serverless-example",
        replace_collection=True,
        embedding=agent_embed_cfg,
    ),
    parsing=ParsingConfig(
        separators=["\n"],
        splitter=Splitter.SIMPLE,
    ),
    n_similar_chunks=2,
    n_relevant_chunks=2,
)

agent = DocChatAgent(config=agent_config)

###################
# Once we have created an agent, we can start loading
# some docs into our Pinecone index:
###################

documents = [
    Document(
        content="""Max Verstappen was the Formula 1 World Drivers' Champion in 2024.
        Lewis Hamilton was the Formula 1 World Drivers' Champion in 2020.
        Nico Rosberg was the Formula 1 World Drivers' Champion in 2016.
        Sebastian Vettel was the Formula 1 World Drivers' Champion in 2013.
        Jenson Button was the Formula 1 World Drivers' Champion in 2009.
        Kimi Räikkönen was the Formula 1 World Drivers' Champion in 2007.
        """,
        metadata=DocMetaData(
            source="wikipedia",
            id="formula-1-facts",
        )
    ),
    Document(
        content="""The Boston Celtics won the NBA Championship for the 2024 NBA season. The MVP for the 2024 NBA Championship was Jaylen Brown.
        The Denver Nuggets won the NBA Championship for the 2023 NBA season. The MVP for the 2023 NBA Championship was Nikola Jokić.
        The Golden State Warriors won the NBA Championship for the 2022 NBA season. The MVP for the 2022 NBA Championship was Stephen Curry.
        The Milwaukee Bucks won the NBA Championship for the 2021 NBA season. The MVP for the 2021 NBA Championship was Giannis Antetokounmpo.
        The Los Angeles Lakers won the NBA Championship for the 2020 NBA season. The MVP for the 2020 NBA Championship was LeBron James.
        The Toronto Raptors won the NBA Championship for the 2019 NBA season. The MVP for the 2019 NBA Championship was Kawhi Leonard.
        """,
        metadata=DocMetaData(
            source="wikipedia",
            id="nba-facts"
        )
    )
]

agent.ingest_docs(documents)

###################
# With the documents now loaded, we can now prompt our agent
###################

formula_one_world_champion_2007 = agent.llm_response(
    message="Who was the Formula 1 World Drivers' Champion in 2007?"
)
try:
    assert "Kimi Räikkönen" in formula_one_world_champion_2007.content
except AssertionError as e:
    print(f"Did not resolve Kimi Räikkönen as the answer, document content: {formula_one_world_champion_2007.content} ")

nba_champion_2023 = agent.llm_response(
    message="Who won the 2023 NBA Championship?"
)
try:
    assert "Denver Nuggets" in nba_champion_2023.content
except AssertionError as e:
    print(f"Did not resolve the Denver Nuggets as the answer, document content: {nba_champion_2023.content}")

nba_mvp_2023 = agent.llm_response(
    message="Who was the MVP for the 2023 NBA Championship?"
)
try:
    assert "Nikola Jokić" in nba_mvp_2023.content
except AssertionError as e:
    print(f"Did not resolve Nikola Jokić as the answer, document content: {nba_mvp_2023.content}")
```
</file>

<file path="docs/notes/portkey.md">
# Portkey Integration

Langroid provides seamless integration with [Portkey](https://portkey.ai), a powerful AI gateway that enables you to access multiple LLM providers through a unified API with advanced features like caching, retries, fallbacks, and comprehensive observability.

## What is Portkey?

Portkey is an AI gateway that sits between your application and various LLM providers, offering:

- **Unified API**: Access 200+ models from different providers through one interface
- **Reliability**: Automatic retries, fallbacks, and load balancing
- **Observability**: Detailed logging, tracing, and analytics
- **Performance**: Intelligent caching and request optimization
- **Security**: Virtual keys and advanced access controls
- **Cost Management**: Usage tracking and budget controls

For complete documentation, visit the [Portkey Documentation](https://docs.portkey.ai).

## Quick Start

### 1. Setup

First, sign up for a Portkey account at [portkey.ai](https://portkey.ai) and get your API key.

Set up your environment variables, either explicitly or in your `.env` file as usual: 

```bash
# Required: Portkey API key
export PORTKEY_API_KEY="your-portkey-api-key"

# Required: Provider API keys (for the models you want to use)
export OPENAI_API_KEY="your-openai-key"
export ANTHROPIC_API_KEY="your-anthropic-key"
export GOOGLE_API_KEY="your-google-key"
# ... other provider keys as needed
```

### 2. Basic Usage

```python
import langroid as lr
import langroid.language_models as lm
from langroid.language_models.provider_params import PortkeyParams

# Create an LLM config to use Portkey's OpenAI-compatible API
# (Note that the name `OpenAIGPTConfig` does NOT imply it only works with OpenAI models;
# the name reflects the fact that the config is meant to be used with an
# OpenAI-compatible API, which Portkey provides for multiple LLM providers.)
llm_config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini",
    portkey_params=PortkeyParams(
        api_key="your-portkey-api-key",  # Or set PORTKEY_API_KEY env var
    )
)

# Create LLM instance
llm = lm.OpenAIGPT(llm_config)

# Use normally
response = llm.chat("What is the smallest prime number?")
print(response.message)
```

### 3. Multiple Providers

Switch between providers seamlessly:

```python
# OpenAI
config_openai = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o",
)

# Anthropic
config_anthropic = lm.OpenAIGPTConfig(
    chat_model="portkey/anthropic/claude-3-5-sonnet-20241022",
)

# Google Gemini
config_gemini = lm.OpenAIGPTConfig(
    chat_model="portkey/google/gemini-2.0-flash-lite",
)
```

## Advanced Features

### Virtual Keys

Use virtual keys to abstract provider management:

```python
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o",
    portkey_params=PortkeyParams(
        virtual_key="vk-your-virtual-key",  # Configured in Portkey dashboard
    )
)
```

### Caching and Performance

Enable intelligent caching to reduce costs and improve performance:

```python
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini",
    portkey_params=PortkeyParams(
        cache={
            "enabled": True,
            "ttl": 3600,  # 1 hour cache
            "namespace": "my-app"
        },
        cache_force_refresh=False,
    )
)
```

### Retry Strategies

Configure automatic retries for better reliability:

```python
config = lm.OpenAIGPTConfig(
    chat_model="portkey/anthropic/claude-3-haiku-20240307",
    portkey_params=PortkeyParams(
        retry={
            "max_retries": 3,
            "backoff": "exponential",
            "jitter": True
        }
    )
)
```

### Observability and Tracing

Add comprehensive tracking for production monitoring:

```python
import uuid

config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o",
    portkey_params=PortkeyParams(
        trace_id=f"trace-{uuid.uuid4().hex[:8]}",
        metadata={
            "user_id": "user-123",
            "session_id": "session-456",
            "app_version": "1.2.3"
        },
        user="user-123",
        organization="my-org",
        custom_headers={
            "x-request-source": "langroid",
            "x-feature": "chat-completion"
        }
    )
)
```

## Configuration Reference

The `PortkeyParams` class supports all Portkey features:

```python
from langroid.language_models.provider_params import PortkeyParams

params = PortkeyParams(
    # Authentication
    api_key="pk-...",                    # Portkey API key
    virtual_key="vk-...",               # Virtual key (optional)
    
    # Observability
    trace_id="trace-123",               # Request tracing
    metadata={"key": "value"},          # Custom metadata
    user="user-id",                     # User identifier
    organization="org-id",              # Organization identifier
    
    # Performance
    cache={                             # Caching configuration
        "enabled": True,
        "ttl": 3600,
        "namespace": "my-app"
    },
    cache_force_refresh=False,          # Force cache refresh
    
    # Reliability
    retry={                             # Retry configuration
        "max_retries": 3,
        "backoff": "exponential",
        "jitter": True
    },
    
    # Custom headers
    custom_headers={                    # Additional headers
        "x-custom": "value"
    },
    
    # Base URL (usually not needed)
    base_url="https://api.portkey.ai"   # Portkey API endpoint
)
```

## Supported Providers

Portkey supports 200+ models from various providers. Common ones include:

```python
# OpenAI
"portkey/openai/gpt-4o"
"portkey/openai/gpt-4o-mini"

# Anthropic
"portkey/anthropic/claude-3-5-sonnet-20241022"
"portkey/anthropic/claude-3-haiku-20240307"

# Google
"portkey/google/gemini-2.0-flash-lite"
"portkey/google/gemini-1.5-pro"

# Cohere
"portkey/cohere/command-r-plus"

# Meta
"portkey/meta/llama-3.1-405b-instruct"

# And many more...
```

Check the [Portkey documentation](https://docs.portkey.ai/docs/integrations/models) for the complete list.

## Examples

Langroid includes comprehensive Portkey examples in `examples/portkey/`:

1. **`portkey_basic_chat.py`** - Basic usage with multiple providers
2. **`portkey_advanced_features.py`** - Caching, retries, and observability
3. **`portkey_multi_provider.py`** - Comparing responses across providers

Run any example:

```bash
cd examples/portkey
python portkey_basic_chat.py
```

## Best Practices

### 1. Use Environment Variables

Never hardcode API keys:

```bash
# .env file
PORTKEY_API_KEY=your_portkey_key
OPENAI_API_KEY=your_openai_key
ANTHROPIC_API_KEY=your_anthropic_key
```

### 2. Implement Fallback Strategies

Use multiple providers for reliability:

```python
providers = [
    ("openai", "gpt-4o-mini"),
    ("anthropic", "claude-3-haiku-20240307"),
    ("google", "gemini-2.0-flash-lite")
]

for provider, model in providers:
    try:
        config = lm.OpenAIGPTConfig(
            chat_model=f"portkey/{provider}/{model}"
        )
        llm = lm.OpenAIGPT(config)
        return llm.chat(question)
    except Exception:
        continue  # Try next provider
```

### 3. Add Meaningful Metadata

Include context for better observability:

```python
params = PortkeyParams(
    metadata={
        "user_id": user.id,
        "feature": "document_qa",
        "document_type": "pdf",
        "processing_stage": "summary"
    }
)
```

### 4. Use Caching Wisely

Enable caching for deterministic queries:

```python
# Good for caching
params = PortkeyParams(
    cache={"enabled": True, "ttl": 3600}
)

# Use with deterministic prompts
response = llm.chat("What is the capital of France?")
```

### 5. Monitor Performance

Use trace IDs to track request flows:

```python
import uuid

trace_id = f"trace-{uuid.uuid4().hex[:8]}"
params = PortkeyParams(
    trace_id=trace_id,
    metadata={"operation": "document_processing"}
)

# Use the same trace_id for related requests
```

## Monitoring and Analytics

### Portkey Dashboard

View detailed analytics at [app.portkey.ai](https://app.portkey.ai):

- Request/response logs
- Token usage and costs
- Performance metrics (latency, errors)
- Provider comparisons
- Custom filters by metadata

### Custom Filtering

Use metadata and headers to filter requests:

```python
# Tag requests by feature
params = PortkeyParams(
    metadata={"feature": "chat", "version": "v2"},
    custom_headers={"x-request-type": "production"}
)
```

Then filter in the dashboard by:
- `metadata.feature = "chat"`
- `headers.x-request-type = "production"`

## Troubleshooting

### Common Issues

1. **Authentication Errors**
   ```
   Error: Unauthorized (401)
   ```
   - Check `PORTKEY_API_KEY` is set correctly
   - Verify API key is active in Portkey dashboard

2. **Provider API Key Missing**
   ```
   Error: Missing API key for provider
   ```
   - Set provider API key (e.g., `OPENAI_API_KEY`)
   - Or use virtual keys in Portkey dashboard

3. **Model Not Found**
   ```
   Error: Model not supported
   ```
   - Check model name format: `portkey/provider/model`
   - Verify model is available through Portkey

4. **Rate Limiting**
   ```
   Error: Rate limit exceeded
   ```
   - Configure retry parameters
   - Use virtual keys for better rate limit management

### Debug Mode

Enable detailed logging:

```python
import logging
logging.getLogger("langroid").setLevel(logging.DEBUG)
```

### Test Configuration

Verify your setup:

```python
# Test basic connection
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini",
    max_output_tokens=50
)
llm = lm.OpenAIGPT(config)
response = llm.chat("Hello")
print("✅ Portkey integration working!")
```

## Migration Guide

### From Direct Provider Access

If you're currently using providers directly:

```python
# Before: Direct OpenAI
config = lm.OpenAIGPTConfig(
    chat_model="gpt-4o-mini"
)

# After: Through Portkey
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini"
)
```

### Adding Advanced Features Gradually

Start simple and add features as needed:

```python
# Step 1: Basic Portkey
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini"
)

# Step 2: Add caching
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini",
    portkey_params=PortkeyParams(
        cache={"enabled": True, "ttl": 3600}
    )
)

# Step 3: Add observability
config = lm.OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini",
    portkey_params=PortkeyParams(
        cache={"enabled": True, "ttl": 3600},
        metadata={"app": "my-app", "user": "user-123"},
        trace_id="trace-abc123"
    )
)
```

## Resources

- **Portkey Website**: [https://portkey.ai](https://portkey.ai)
- **Portkey Documentation**: [https://docs.portkey.ai](https://docs.portkey.ai)
- **Portkey Dashboard**: [https://app.portkey.ai](https://app.portkey.ai)
- **Supported Models**: [https://docs.portkey.ai/docs/integrations/models](https://docs.portkey.ai/docs/integrations/models)
- **Langroid Examples**: `examples/portkey/` directory
- **API Reference**: [https://docs.portkey.ai/docs/api-reference](https://docs.portkey.ai/docs/api-reference)
</file>

<file path="docs/notes/pydantic-v2-migration.md">
# Pydantic v2 Migration Guide

## Overview

Langroid has fully migrated to Pydantic v2! All internal code now uses Pydantic v2 
patterns and imports directly from `pydantic`. This guide will help you update your 
code to work with the new version.

## Compatibility Layer (Deprecated)

If your code currently imports from `langroid.pydantic_v1`:

```python
# OLD - Deprecated
from langroid.pydantic_v1 import BaseModel, Field, BaseSettings
```

You'll see a deprecation warning. This compatibility layer now imports from Pydantic v2 
directly, so your code may continue to work, but you should update your imports:

```python
# NEW - Correct
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings  # Note: BaseSettings moved to pydantic_settings in v2
```

!!! note "BaseSettings Location Change"
    In Pydantic v2, `BaseSettings` has moved to a separate `pydantic_settings` package.
    You'll need to install it separately: `pip install pydantic-settings`

!!! warning "Compatibility Layer Removal"
    The `langroid.pydantic_v1` module will be removed in a future version. 
    Update your imports now to avoid breaking changes.

## Key Changes to Update

### 1. All Fields Must Have Type Annotations

!!! danger "Critical Change"
    In Pydantic v2, fields without type annotations are completely ignored!

```python
# WRONG - Fields without annotations are ignored in v2
class MyModel(BaseModel):
    name = "John"          # ❌ This field is IGNORED!
    age = 25               # ❌ This field is IGNORED!
    role: str = "user"     # ✅ This field works

# CORRECT - All fields must have type annotations
class MyModel(BaseModel):
    name: str = "John"     # ✅ Type annotation required
    age: int = 25          # ✅ Type annotation required
    role: str = "user"     # ✅ Already correct
```

This is one of the most common issues when migrating to v2. Always ensure every field has an explicit type annotation, even if it has a default value.

#### Special Case: Overriding Fields in Subclasses

!!! danger "Can Cause Errors"
    When overriding fields from parent classes without type annotations, you may get 
    actual errors, not just ignored fields!

This is particularly important when creating custom Langroid agent configurations:

```python
# WRONG - This can cause errors!
from langroid import ChatAgentConfig
from langroid.language_models import OpenAIGPTConfig

class MyAgentConfig(ChatAgentConfig):
    # ❌ ERROR: Missing type annotation when overriding parent field
    llm = OpenAIGPTConfig(chat_model="gpt-4")
    
    # ❌ ERROR: Even with Field, still needs type annotation
    system_message = Field(default="You are a helpful assistant")

# CORRECT - Always include type annotations when overriding
class MyAgentConfig(ChatAgentConfig):
    # ✅ Type annotation required when overriding
    llm: OpenAIGPTConfig = OpenAIGPTConfig(chat_model="gpt-4")
    
    # ✅ Type annotation with Field
    system_message: str = Field(default="You are a helpful assistant")
```

Without type annotations on overridden fields, you may see errors like:
- `ValueError: Field 'llm' requires a type annotation`
- `TypeError: Field definitions should be annotated`
- Validation errors when the model tries to use the parent's field definition

### 2. Stricter Type Validation for Optional Fields

!!! danger "Breaking Change"
    Pydantic v2 is much stricter about type validation. Fields that could accept `None` 
    in v1 now require explicit `Optional` type annotations.

```python
# WRONG - This worked in v1 but fails in v2
class CloudSettings(BaseSettings):
    private_key: str = None      # ❌ ValidationError: expects string, got None
    api_host: str = None         # ❌ ValidationError: expects string, got None

# CORRECT - Explicitly mark fields as optional
from typing import Optional

class CloudSettings(BaseSettings):
    private_key: Optional[str] = None    # ✅ Explicitly optional
    api_host: Optional[str] = None       # ✅ Explicitly optional
    
    # Or using Python 3.10+ union syntax
    client_email: str | None = None      # ✅ Also works
```

This commonly affects:
- Configuration classes using `BaseSettings`
- Fields with `None` as default value
- Environment variable loading where the var might not be set

If you see errors like:
```
ValidationError: Input should be a valid string [type=string_type, input_value=None, input_type=NoneType]
```

The fix is to add `Optional[]` or `| None` to the type annotation.

### 3. Model Serialization Methods

```python
# OLD (Pydantic v1)
data = model.dict()
json_str = model.json()
new_model = MyModel.parse_obj(data)
new_model = MyModel.parse_raw(json_str)

# NEW (Pydantic v2)
data = model.model_dump()
json_str = model.model_dump_json()
new_model = MyModel.model_validate(data)
new_model = MyModel.model_validate_json(json_str)
```

### 4. Model Configuration

```python
# OLD (Pydantic v1)
class MyModel(BaseModel):
    name: str
    
    class Config:
        extra = "forbid"
        validate_assignment = True

# NEW (Pydantic v2)
from pydantic import BaseModel, ConfigDict

class MyModel(BaseModel):
    model_config = ConfigDict(
        extra="forbid",
        validate_assignment=True
    )
    
    name: str
```

### 5. Field Validators

```python
# OLD (Pydantic v1)
from pydantic import validator

class MyModel(BaseModel):
    name: str
    
    @validator('name')
    def name_must_not_be_empty(cls, v):
        if not v.strip():
            raise ValueError('Name cannot be empty')
        return v

# NEW (Pydantic v2)
from pydantic import field_validator

class MyModel(BaseModel):
    name: str
    
    @field_validator('name')
    def name_must_not_be_empty(cls, v):
        if not v.strip():
            raise ValueError('Name cannot be empty')
        return v
```

### 6. Custom Types and Validation

```python
# OLD (Pydantic v1)
from pydantic import parse_obj_as
from typing import List

data = [{"name": "Alice"}, {"name": "Bob"}]
users = parse_obj_as(List[User], data)

# NEW (Pydantic v2)
from pydantic import TypeAdapter
from typing import List

data = [{"name": "Alice"}, {"name": "Bob"}]
users = TypeAdapter(List[User]).validate_python(data)
```

## Common Patterns in Langroid

When working with Langroid's agents and tools:

### Tool Messages

```python
from pydantic import BaseModel, Field
from langroid.agent.tool_message import ToolMessage

class MyTool(ToolMessage):
    request: str = "my_tool"
    purpose: str = "Process some data"
    
    # Use Pydantic v2 patterns
    data: str = Field(..., description="The data to process")
    
    def handle(self) -> str:
        # Tool logic here
        return f"Processed: {self.data}"
```

### Agent Configuration

```python
from pydantic import ConfigDict
from langroid import ChatAgentConfig

class MyAgentConfig(ChatAgentConfig):
    model_config = ConfigDict(extra="forbid")
    
    custom_param: str = "default_value"
```

## Troubleshooting

### Import Errors

If you see `ImportError` or `AttributeError` after updating imports:
- Make sure you're using the correct v2 method names (e.g., `model_dump` not `dict`)
- Check that field validators use `@field_validator` not `@validator`
- Ensure `ConfigDict` is used instead of nested `Config` classes

### Validation Errors

Pydantic v2 has stricter validation in some cases:
- Empty strings are no longer coerced to `None` for optional fields
- Type coercion is more explicit
- Extra fields handling may be different

### Performance

Pydantic v2 is generally faster, but if you notice any performance issues:
- Use `model_validate` instead of creating models with `**dict` unpacking
- Consider using `model_construct` for trusted data (skips validation)

## Need Help?

If you encounter issues during migration:
1. Check the [official Pydantic v2 migration guide](https://docs.pydantic.dev/latest/migration/)
2. Review Langroid's example code for v2 patterns
3. Open an issue on the [Langroid GitHub repository](https://github.com/langroid/langroid/issues)
</file>

<file path="docs/notes/qdrant-resource-cleanup.md">
# QdrantDB Resource Cleanup

When using QdrantDB with local storage, it's important to properly release resources
to avoid file lock conflicts. QdrantDB uses a `.lock` file to prevent concurrent
access to the same storage directory.

## The Problem

Without proper cleanup, you may encounter this warning:

```
Error connecting to local QdrantDB at ./qdrant_data:
Storage folder ./qdrant_data is already accessed by another instance of Qdrant
client. If you require concurrent access, use Qdrant server instead.
Switching to ./qdrant_data.new
```

This happens when a QdrantDB instance isn't properly closed, leaving the lock file
in place.

## Solutions

### Method 1: Explicit `close()` Method

Always call `close()` when done with a QdrantDB instance:

```python
from langroid.vector_store.qdrantdb import QdrantDB, QdrantDBConfig

config = QdrantDBConfig(
    cloud=False,
    collection_name="my_collection",
    storage_path="./qdrant_data",
)

vecdb = QdrantDB(config)
# ... use the vector database ...
vecdb.clear_all_collections(really=True)

# Important: Release the lock
vecdb.close()
```

### Method 2: Context Manager (Recommended)

Use QdrantDB as a context manager for automatic cleanup:

```python
from langroid.vector_store.qdrantdb import QdrantDB, QdrantDBConfig

config = QdrantDBConfig(
    cloud=False,
    collection_name="my_collection", 
    storage_path="./qdrant_data",
)

with QdrantDB(config) as vecdb:
    # ... use the vector database ...
    vecdb.clear_all_collections(really=True)
    # Automatically closed when exiting the context
```

The context manager ensures cleanup even if an exception occurs.

## When This Matters

This is especially important in scenarios where:

1. You create temporary QdrantDB instances for maintenance (e.g., clearing
   collections)
2. Your application restarts frequently during development
3. Multiple parts of your code need to access the same storage path sequentially

## Note for Cloud Storage

This only affects local storage (`cloud=False`). When using Qdrant cloud service,
the lock file mechanism is not used.
</file>

<file path="docs/notes/quiet-mode.md">
# Suppressing LLM output: quiet mode

In some scenarios we want to suppress LLM streaming output -- e.g. when doing some type of processing as part of a workflow,
or when using an LLM-agent to generate code via tools, etc. We are more interested in seeing the results of the workflow,
and don't want to see streaming output in the terminal. Langroid provides a `quiet_mode` context manager that can be used
to suppress LLM output, even in streaming mode (in fact streaming is disabled in quiet mode).

E.g.  we can use the `quiet_mode` context manager like this:

```python
from langroid.utils.configuration import quiet_mode, settings

# directly with LLM

llm = ...
with quiet_mode(True):
	response = llm.chat(...)

# or, using an agent

agent = ...
with quiet_mode(True):
	response = agent.llm_response(...)

# or, using a task

task = Task(agent, ...)
with quiet_mode(True):
	result = Taks.run(...)

# we can explicitly set quiet_mode, and this is globally recognized throughout langroid.

settings.quiet = True

# we can also condition quiet mode on another custom cmd line option/flag, such as "silent":

with quiet_mode(silent):
	...

```
</file>

<file path="docs/notes/structured-output.md">
# Structured Output

Available in Langroid since v0.24.0.

On supported LLMs, including recent OpenAI LLMs (GPT-4o and GPT-4o mini) and local LLMs served by compatible inference servers,
in particular, [vLLM](https://github.com/vllm-project/vllm) and [llama.cpp](https://github.com/ggerganov/llama.cpp), the decoding process can be constrained to ensure that the model's output adheres to a provided schema, 
improving the reliability of tool call generation and, in general, ensuring that the output can be reliably parsed and processed by downstream applications.

See [here](../tutorials/local-llm-setup.md/#setup-llamacpp-with-a-gguf-model-from-huggingface) for instructions for usage with `llama.cpp` and [here](../tutorials/local-llm-setup.md/#setup-vllm-with-a-model-from-huggingface) for `vLLM`.

Given a `ChatAgent` `agent` and a type `type`, we can define a strict copy of the agent as follows:
```python
strict_agent = agent[type]
```

We can use this to allow reliable extraction of typed values from an LLM with minimal prompting. For example, to generate typed values given `agent`'s current context, we can define the following:

```python
def typed_agent_response(
    prompt: str,
    output_type: type,
) -> Any:
    response = agent[output_type].llm_response_forget(prompt)
    return agent.from_ChatDocument(response, output_type)
```

We apply this in [test_structured_output.py](https://github.com/langroid/langroid/blob/main/tests/main/test_structured_output.py), in which we define types which describe
countries and their presidents:
```python
class Country(BaseModel):
    """Info about a country"""

    name: str = Field(..., description="Name of the country")
    capital: str = Field(..., description="Capital of the country")


class President(BaseModel):
    """Info about a president of a country"""

    country: Country = Field(..., description="Country of the president")
    name: str = Field(..., description="Name of the president")
    election_year: int = Field(..., description="Year of election of the president")


class PresidentList(BaseModel):
    """List of presidents of various countries"""

    presidents: List[President] = Field(..., description="List of presidents")
```
and show that `typed_agent_response("Show me an example of two Presidents", PresidentsList)` correctly returns a list of two presidents with *no* prompting describing the desired output format.

In addition to Pydantic models, `ToolMessage`s, and simple Python types are supported. For instance, `typed_agent_response("What is the value of pi?", float)` correctly returns $\pi$ to several decimal places.

The following two detailed examples show how structured output can be used to improve the reliability of the [chat-tree example](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tree.py): [this](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tree-structured.py) shows how we can use output formats to force the agent to make the correct tool call in each situation and [this](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tree-structured-simple.py) shows how we can simplify by using structured outputs to extract typed intermediate values and expressing the control flow between LLM calls and agents explicitly.
</file>

<file path="docs/notes/task-termination.md">
# Task Termination in Langroid

## Why Task Termination Matters

When building agent-based systems, one of the most critical yet challenging aspects is determining when a task should complete. Unlike traditional programs with clear exit points, agent conversations can meander, loop, or continue indefinitely. Getting termination wrong leads to two equally problematic scenarios:

**Terminating too early** means missing crucial information or cutting off an agent mid-process. Imagine an agent that searches for information, finds it, but terminates before it can process or summarize the results. The task completes "successfully" but fails to deliver value.

**Terminating too late** wastes computational resources, frustrates users, and can lead to repetitive loops where agents keep responding without making progress. We've all experienced chatbots that won't stop talking or systems that keep asking "Is there anything else?" long after the conversation should have ended. Even worse, agents can fall into infinite loops—repeatedly exchanging the same messages, calling the same tools, or cycling through states without making progress. These loops not only waste resources but can rack up significant costs when using paid LLM APIs.

The challenge is that the "right" termination point depends entirely on context. A customer service task might complete after resolving an issue and confirming satisfaction. A research task might need to gather multiple sources, synthesize them, and present findings. A calculation task should end after computing and presenting the result. Each scenario requires different termination logic.

Traditionally, developers would subclass `Task` and override the `done()` method with custom logic. While flexible, this approach scattered termination logic across multiple subclasses, making systems harder to understand and maintain. It also meant that common patterns—like "complete after tool use" or "stop when the user says goodbye"—had to be reimplemented repeatedly.

This guide introduces Langroid's declarative approach to task termination, culminating in the powerful `done_sequences` feature. Instead of writing imperative code, you can now describe *what* patterns should trigger completion, and Langroid handles the *how*. This makes your agent systems more predictable, maintainable, and easier to reason about.

## Table of Contents
- [Overview](#overview)
- [Basic Termination Methods](#basic-termination-methods)
- [Done Sequences: Event-Based Termination](#done-sequences-event-based-termination)
  - [Concept](#concept)
  - [DSL Syntax (Recommended)](#dsl-syntax-recommended)
  - [Full Object Syntax](#full-object-syntax)
  - [Event Types](#event-types)
  - [Examples](#examples)
- [Implementation Details](#implementation-details)
- [Best Practices](#best-practices)
- [Reference](#reference)
- [Text-Based Termination Signals](#text-based-termination-signals)

## Overview

In Langroid, a `Task` wraps an `Agent` and manages the conversation flow. Controlling when a task terminates is crucial for building reliable agent systems. Langroid provides several methods for task termination, from simple flags to sophisticated event sequence matching.

## Basic Termination Methods

### 1. Turn Limits
```python
# Task runs for exactly 5 turns
result = task.run("Start conversation", turns=5)
```

### 2. Single Round Mode
```python
# Task completes after one exchange
config = TaskConfig(single_round=True)
task = Task(agent, config=config)
```

### 3. Done If Tool
```python
# Task completes when any tool is generated
config = TaskConfig(done_if_tool=True)
task = Task(agent, config=config)
```

### 4. Done If Response/No Response
```python
# Task completes based on response from specific entities
config = TaskConfig(
    done_if_response=[Entity.LLM],      # Done if LLM responds
    done_if_no_response=[Entity.USER]   # Done if USER doesn't respond
)
```

### 5. String Signals
```python
# Task completes when special strings like "DONE" are detected
# (enabled by default with recognize_string_signals=True)
```

See [Text-Based Routing and Signal Control](#text-based-routing-and-signal-control)
for detailed documentation on controlling text-based routing behavior.

### 6. Orchestration Tools
```python
# Using DoneTool, FinalResultTool, etc.
from langroid.agent.tools.orchestration import DoneTool
agent.enable_message(DoneTool)
```

## Done Sequences: Event-Based Termination

### Concept

The `done_sequences` feature allows you to specify sequences of events that trigger task completion. This provides fine-grained control over task termination based on conversation patterns.

**Key Features:**

- Specify multiple termination sequences
- Use convenient DSL syntax or full object syntax
- Strict consecutive matching (no skipping events)
- Efficient implementation using message parent pointers

### DSL Syntax (Recommended)

The DSL (Domain Specific Language) provides a concise way to specify sequences:

```python
from langroid.agent.task import Task, TaskConfig

config = TaskConfig(
    done_sequences=[
        "T, A",                    # Tool followed by agent response
        "T[calculator], A",        # Specific calculator tool by name
        "T[CalculatorTool], A",    # Specific tool by class reference (NEW!)
        "L, T, A, L",              # LLM, tool, agent, LLM sequence
        "C[quit|exit|bye]",        # Content matching regex
        "U, L, A",                 # User, LLM, agent sequence
    ]
)
task = Task(agent, config=config)
```

#### DSL Pattern Reference

| Pattern | Description | Event Type |
|---------|-------------|------------|
| `T` | Any tool | `TOOL` |
| `T[name]` | Specific tool by name | `SPECIFIC_TOOL` |
| `T[ToolClass]` | Specific tool by class (NEW!) | `SPECIFIC_TOOL` |
| `A` | Agent response | `AGENT_RESPONSE` |
| `L` | LLM response | `LLM_RESPONSE` |
| `U` | User response | `USER_RESPONSE` |
| `N` | No response | `NO_RESPONSE` |
| `C[pattern]` | Content matching regex | `CONTENT_MATCH` |

**Examples:**

- `"T, A"` - Any tool followed by agent handling
- `"T[search], A, T[calculator], A"` - Search tool, then calculator tool
- `"T[CalculatorTool], A"` - Specific tool class followed by agent handling (NEW!)
- `"L, C[complete|done|finished]"` - LLM response containing completion words
- `"TOOL, AGENT"` - Full words also supported

### Full Object Syntax

For more control, use the full object syntax:

```python
from langroid.agent.task import (
    Task, TaskConfig, DoneSequence, AgentEvent, EventType
)

config = TaskConfig(
    done_sequences=[
        DoneSequence(
            name="tool_handled",
            events=[
                AgentEvent(event_type=EventType.TOOL),
                AgentEvent(event_type=EventType.AGENT_RESPONSE),
            ]
        ),
        DoneSequence(
            name="specific_tool_pattern",
            events=[
                AgentEvent(
                    event_type=EventType.SPECIFIC_TOOL,
                    tool_name="calculator",
                    # Can also use tool_class for type-safe references (NEW!):
                    # tool_class=CalculatorTool
                ),
                AgentEvent(event_type=EventType.AGENT_RESPONSE),
            ]
        ),
    ]
)
```

### Event Types

The following event types are available:

| EventType | Description | Additional Parameters |
|-----------|-------------|----------------------|
| `TOOL` | Any tool message generated | - |
| `SPECIFIC_TOOL` | Specific tool by name or class | `tool_name`, `tool_class` (NEW!) |
| `LLM_RESPONSE` | LLM generates a response | - |
| `AGENT_RESPONSE` | Agent responds (e.g., handles tool) | - |
| `USER_RESPONSE` | User provides input | - |
| `CONTENT_MATCH` | Response matches regex pattern | `content_pattern` |
| `NO_RESPONSE` | No valid response from entity | - |

### Examples

#### Example 1: Tool Completion
Task completes after any tool is used and handled:

```python
config = TaskConfig(done_sequences=["T, A"])
```

This is equivalent to `done_if_tool=True` but happens after the agent handles the tool.

#### Example 2: Multi-Step Process
Task completes after a specific conversation pattern:

```python
config = TaskConfig(
    done_sequences=["L, T[calculator], A, L"]
)
# Completes after: LLM response → calculator tool → agent handles → LLM summary
```

#### Example 3: Multiple Exit Conditions
Different ways to complete the task:

```python
config = TaskConfig(
    done_sequences=[
        "C[quit|exit|bye]",           # User says quit
        "T[calculator], A",           # Calculator used
        "T[search], A, T[search], A", # Two searches performed
    ]
)
```

#### Example 4: Tool Class References (NEW!)
Use actual tool classes instead of string names for type safety:

```python
from langroid.agent.tool_message import ToolMessage

class CalculatorTool(ToolMessage):
    request: str = "calculator"
    # ... tool implementation

class SearchTool(ToolMessage):
    request: str = "search"
    # ... tool implementation

# Enable tools on the agent
agent.enable_message([CalculatorTool, SearchTool])

# Use tool classes in done sequences
config = TaskConfig(
    done_sequences=[
        "T[CalculatorTool], A",  # Using class name
        "T[SearchTool], A, T[CalculatorTool], A",  # Multiple tools
    ]
)
```

**Benefits of tool class references:**
- **Type-safe**: IDE can validate tool class names
- **Refactoring-friendly**: Renaming tool classes automatically updates references
- **No string typos**: Compiler/linter catches invalid class names
- **Better IDE support**: Autocomplete and go-to-definition work

#### Example 5: Mixed Syntax
Combine DSL strings and full objects:

```python
config = TaskConfig(
    done_sequences=[
        "T, A",  # Simple DSL
        "T[CalculatorTool], A",  # Tool class reference (NEW!)
        DoneSequence(  # Full control
            name="complex_check",
            events=[
                AgentEvent(
                    event_type=EventType.SPECIFIC_TOOL,
                    tool_name="database_query",
                    tool_class=DatabaseQueryTool,  # Can use class directly (NEW!)
                    responder="DatabaseAgent"
                ),
                AgentEvent(event_type=EventType.AGENT_RESPONSE),
            ]
        ),
    ]
)
```

## Implementation Details

### How Done Sequences Work

Done sequences operate at the **task level** and are based on the **sequence of valid responses** generated during a task's execution. When a task runs, it maintains a `response_sequence` that tracks each message (ChatDocument) as it's processed.

**Key points:**
- Done sequences are checked only within a single task's scope
- They track the temporal order of responses within that task
- The response sequence is built incrementally as the task processes each step
- Only messages that represent valid responses are added to the sequence

### Response Sequence Building
The task builds its response sequence during execution:

```python
# In task.run(), after each step:
if self.pending_message is not None:
    if (not self.response_sequence or 
        self.pending_message.id() != self.response_sequence[-1].id()):
        self.response_sequence.append(self.pending_message)
```

### Message Chain Retrieval
Done sequences are checked against the response sequence:

```python
def _get_message_chain(self, msg: ChatDocument, max_depth: Optional[int] = None):
    """Get the chain of messages from response sequence"""
    if max_depth is None:
        max_depth = 50  # default
        if self._parsed_done_sequences:
            max_depth = max(len(seq.events) for seq in self._parsed_done_sequences)
    
    # Simply return the last max_depth elements from response_sequence
    return self.response_sequence[-max_depth:]
```

**Note:** The response sequence used for done sequences is separate from the parent-child pointer system. Parent pointers track causal relationships and lineage across agent boundaries (important for debugging and understanding delegation patterns), while response sequences track temporal order within a single task for termination checking.

### Strict Matching
Events must occur consecutively without intervening messages:

```python
# This sequence: [TOOL, AGENT_RESPONSE]
# Matches: USER → LLM(tool) → AGENT
# Does NOT match: USER → LLM(tool) → USER → AGENT
```

### Performance

- Efficient O(n) traversal where n is sequence length
- No full history scan needed
- Early termination on first matching sequence

## Best Practices

1. **Use DSL for Simple Cases**
   ```python
   # Good: Clear and concise
   done_sequences=["T, A"]
   
   # Avoid: Verbose for simple patterns
   done_sequences=[DoneSequence(events=[...])]
   ```

2. **Name Your Sequences**
   ```python
   DoneSequence(
       name="calculation_complete",  # Helps with debugging
       events=[...]
   )
   ```

3. **Order Matters**
   - Put more specific sequences first
   - General patterns at the end

4. **Test Your Sequences**
   ```python
   # Use MockLM for testing
   agent = ChatAgent(
       ChatAgentConfig(
           llm=MockLMConfig(response_fn=lambda x: "test response")
       )
   )
   ```

5. **Combine with Other Methods**
   ```python
   config = TaskConfig(
       done_if_tool=True,      # Quick exit on any tool
       done_sequences=["L, L, L"],  # Or after 3 LLM responses
       max_turns=10,           # Hard limit
   )
   ```

## Reference

### Code Examples
- **Basic example**: [`examples/basic/done_sequences_example.py`](../../examples/basic/done_sequences_example.py)
- **Test cases**: [`tests/main/test_done_sequences.py`](../../tests/main/test_done_sequences.py) (includes tool class tests)
- **DSL tests**: [`tests/main/test_done_sequences_dsl.py`](../../tests/main/test_done_sequences_dsl.py)
- **Parser tests**: [`tests/main/test_done_sequence_parser.py`](../../tests/main/test_done_sequence_parser.py)

### Core Classes
- `TaskConfig` - Configuration including `done_sequences`
- `DoneSequence` - Container for event sequences
- `AgentEvent` - Individual event in a sequence
- `EventType` - Enumeration of event types

### Parser Module
- `langroid.agent.done_sequence_parser` - DSL parsing functionality

### Task Methods
- `Task.done()` - Main method that checks sequences
- `Task._matches_sequence_with_current()` - Sequence matching logic
- `Task._classify_event()` - Event classification
- `Task._get_message_chain()` - Message traversal

## Migration Guide

If you're currently overriding `Task.done()`:

```python
# Before: Custom done() method
class MyTask(Task):
    def done(self, result=None, r=None):
        if some_complex_logic(result):
            return (True, StatusCode.DONE)
        return super().done(result, r)

# After: Use done_sequences
config = TaskConfig(
    done_sequences=["T[my_tool], A, L"]  # Express as sequence
)
task = Task(agent, config=config)  # No subclassing needed
```

**NEW: Using Tool Classes Instead of Strings**

If you have tool classes defined, you can now reference them directly:

```python
# Before: Using string names (still works)
config = TaskConfig(
    done_sequences=["T[calculator], A"]  # String name
)

# After: Using tool class references (recommended)
config = TaskConfig(
    done_sequences=["T[CalculatorTool], A"]  # Class name
)
```

This provides better type safety and makes refactoring easier.

## Troubleshooting

**Sequence not matching?**

- Check that events are truly consecutive (no intervening messages)
- Use logging to see the actual message chain
- Verify tool names match exactly

**Type errors with DSL?**

- Ensure you're using strings for DSL patterns
- Check that tool names in `T[name]` don't contain special characters

**Performance concerns?**

- Sequences only traverse as deep as needed
- Consider shorter sequences for better performance
- Use specific tool names to avoid unnecessary checks

## Text-Based Termination Signals

### `TaskConfig.recognize_string_signals`

Controls whether the task recognizes text-based orchestration signals like `DONE`,
`PASS`, `DONE_PASS`, etc.

```python
from langroid.agent.task import Task, TaskConfig

# Default: signals are recognized
task = Task(agent, config=TaskConfig(recognize_string_signals=True))

# Disable: signals treated as plain text
task = Task(agent, config=TaskConfig(recognize_string_signals=False))
```

**When `True` (default):**

- `DONE` in a response signals task completion
- `PASS` signals passing control to another agent
- `DONE_PASS` combines both behaviors

**When `False`:**

- These strings are treated as literal text
- Useful when LLM responses might accidentally contain these keywords
- Task termination must use other mechanisms (tools, `done_sequences`, etc.)

Note that `PASS` also relates to message routing between agents. For more details
on text-based routing and the related `recognize_recipient_in_content` setting,
see [Message Routing](message-routing.md).

## Summary

The `done_sequences` feature provides a powerful, declarative way to control task
termination based on conversation patterns. The DSL syntax makes common cases
simple while the full object syntax provides complete control when needed. This
approach eliminates the need to subclass `Task` and override `done()` for most
use cases, leading to cleaner, more maintainable code.
</file>

<file path="docs/notes/task-tool.md">
# TaskTool: Spawning Sub-Agents for Task Delegation

## Overview

`TaskTool` allows agents to **spawn sub-agents** to handle specific tasks. When an agent encounters a task that requires specialized tools or isolated execution, it can spawn a new sub-agent with exactly the capabilities needed for that task.

This enables agents to dynamically create a hierarchy of specialized workers, each focused on their specific subtask with only the tools they need.

## When to Use TaskTool

TaskTool is useful when:
- Different parts of a task require different specialized tools
- You want to isolate tool access for specific operations  
- A task involves recursive or nested operations
- You need different LLM models for different subtasks

## How It Works

1. The parent agent decides to spawn a sub-agent and specifies:
   - A system message defining the sub-agent's role
   - A prompt for the sub-agent to process
   - Which tools the sub-agent should have access to
   - Optional model and iteration limits

2. TaskTool spawns the new sub-agent, runs the task, and returns the result to the parent.

## Async Support

TaskTool fully supports both synchronous and asynchronous execution. The tool automatically handles async contexts when the parent task is running asynchronously.

## Usage Example

```python
from langroid.agent.tools.task_tool import TaskTool

# Enable TaskTool for your agent
agent.enable_message([TaskTool, YourCustomTool], use=True, handle=True)

# Agent can now spawn sub-agents for tasks when the LLM generates a task_tool request:

response = {
    "request": "task_tool",
    "system_message": "You are a calculator. Use the multiply_tool to compute products.",
    "prompt": "Calculate 5 * 7",
    "tools": ["multiply_tool"],
    "model": "gpt-4o-mini",   # optional
    "max_iterations": 5,      # optional
    "agent_name": "calculator-agent"  # optional
}
```

## Field Reference

**Required fields:**
- `system_message`: Instructions for the sub-agent's role and behavior
- `prompt`: The specific task/question for the sub-agent
- `tools`: List of tool names. Special values: `["ALL"]` or `["NONE"]`

**Optional fields:**
- `model`: LLM model name (default: "gpt-4o-mini")
- `max_iterations`: Task iteration limit (default: 10)
- `agent_name`: Name for the sub-agent (default: auto-generated as "agent-{uuid}")

## Example: Nested Operations

Consider computing `Nebrowski(10, Nebrowski(3, 2))` where Nebrowski is a custom operation. The main agent spawns sub-agents to handle each operation:

```python
# Main agent spawns first sub-agent for inner operation:
{
    "request": "task_tool",
    "system_message": "Compute Nebrowski operations using the nebrowski_tool.",
    "prompt": "Compute Nebrowski(3, 2)",
    "tools": ["nebrowski_tool"]
}

# Then spawns another sub-agent for outer operation:
{
    "request": "task_tool",
    "system_message": "Compute Nebrowski operations using the nebrowski_tool.",
    "prompt": "Compute Nebrowski(10, 11)",  # where 11 is the previous result
    "tools": ["nebrowski_tool"]
}
```

## Working Examples

See [`tests/main/test_task_tool.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_task_tool.py) for complete examples including:
- Basic task delegation with mock agents
- Nested operations with custom tools
- Both sync and async usage patterns

## Important Notes

- Spawned sub-agents run non-interactively (no human input)
- `DoneTool` is automatically enabled for all sub-agents
- Results are returned as `ChatDocument` objects. The Langroid framework takes care
  of converting them to a suitable format for the parent agent's LLM to consume and 
  respond to.
- Sub-agents can be given custom names via the `agent_name` parameter, which helps with 
  logging and debugging. If not specified, a unique name is auto-generated in the format 
  "agent-{uuid}"
- Only tools "known" to the parent agent can be enabled for sub-agents. This is an 
  important aspect of the current mechanism. The `TaskTool` handler method in
  the sub-agent only has access to tools that are known to the parent agent.
  If there are tools that are only relevant to the sub-agent but not the parent,
  you must still enable them in the parent agent, but you can set `use=False`
  and `handle=False` when you enable them, e.g.:

```python
agent.enable_message(MySubAgentTool, use=False, handle=False)
```
  Since we are letting the main agent's LLM "decide" when to spawn a sub-agent,
  your system message of the main agent should contain instructions clarifying that
  it can decide which tools to enable for the sub-agent, as well as a list of 
  all tools that might possibly be relevant to the sub-agent. This is particularly
  important for tools that have been enabled with `use=False`, since instructions for
  such tools would not be auto-inserted into the agent's system message. 



## Best Practices

1. **Clear Instructions**: Provide specific system messages that explain the sub-agent's role and tool usage
2. **Tool Availability**: Ensure delegated tools are enabled for the parent agent
3. **Appropriate Models**: Use simpler/faster models for simple subtasks
4. **Iteration Limits**: Set reasonable limits based on task complexity
</file>

<file path="docs/notes/tavily_search.md">
---

# **Using Tavily Search with Langroid**

---

## **1. Set Up Tavily**

1. **Access Tavily Platform**  
   Go to the [Tavily Platform](https://tavily.com/).
   
2. **Sign Up or Log In**  
   Create an account or log in if you already have one.

3. **Get Your API Key**  
   - Navigate to your dashboard
   - Copy your API key

4. **Set Environment Variable**  
   Add the following variable to your `.env` file:
   ```env
   TAVILY_API_KEY=<your_api_key>

---

## **2. Use Tavily Search with Langroid**

### **Installation**

```bash
uv add tavily-python
# or
pip install tavily-python
```
### **Code Example**

```python
import langroid as lr
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.tools.tavily_search_tool import TavilySearchTool

# Configure the ChatAgent
config = ChatAgentConfig(
    name="search-agent",
    llm=lr.language_models.OpenAIGPTConfig(
        chat_model=lr.language_models.OpenAIChatModel.GPT4o
    ),
    use_tools=True
)

# Create the agent
agent = ChatAgent(config)

# Enable Tavily search tool
agent.enable_message(TavilySearchTool)

```
---

## **3. Perform Web Searches**

Use the agent to perform web searches using Tavily's AI-powered search.

```python
# Simple search query
response = agent.llm_response(
    "What are the latest developments in quantum computing?"
)
print(response)

# Search with specific number of results
response = agent.llm_response(
    "Find 5 recent news articles about artificial intelligence."
)
print(response)
```
---

## **4. Custom Search Requests**

You can also customize the search behavior by creating a TavilySearchTool instance directly:

```python
from langroid.agent.tools.tavily_search_tool import TavilySearchTool

# Create a custom search request
search_request = TavilySearchTool(
    query="Latest breakthroughs in fusion energy",
    num_results=3
)

# Get search results
results = search_request.handle()
print(results)
```

---
</file>

<file path="docs/notes/tool-message-handler.md">
# Tool Message Handlers in Langroid

## Overview

Langroid provides flexible ways to define handlers for `ToolMessage` classes. When a tool is used by an LLM, the framework needs to know how to handle it. This can be done either by defining a handler method in the `Agent` class or within the `ToolMessage` class itself.

## Enabling Tools with `enable_message`

Before an agent can use or handle a tool, it must be explicitly enabled using the `enable_message` method. This method takes two important arguments:

- **`use`** (bool): Whether the LLM is allowed to generate this tool
- **`handle`** (bool): Whether the agent is allowed to handle this tool

```python
# Enable both generation and handling (default)
agent.enable_message(MyTool, use=True, handle=True)

# Enable only handling (agent can handle but LLM won't generate)
agent.enable_message(MyTool, use=False, handle=True)

# Enable only generation (LLM can generate but agent won't handle)
agent.enable_message(MyTool, use=True, handle=False)
```

When `handle=True` and the `ToolMessage` has a `handle` method defined, this method is inserted into the agent with a name matching the tool's `request` field value. This insertion only happens when `enable_message` is called.

## Default Handler Mechanism

By default, `ToolMessage` uses and/or creates a handler in `Agent` class instance with the name identical to the tool's `request` attribute.

### Agent-based Handlers
If a tool `MyTool` has `request` attribute `my_tool`, you can define a method `my_tool` in your `Agent` class that will handle this tool when the LLM generates it:

```python
class MyTool(ToolMessage):
    request = "my_tool"
    param: str

class MyAgent(ChatAgent):
    def my_tool(self, msg: MyTool) -> str:
        return f"Handled: {msg.param}"

# Enable the tool
agent = MyAgent()
agent.enable_message(MyTool)
```

### ToolMessage-based Handlers
Alternatively, if a tool is "stateless" (i.e. does not require the Agent's state), you can define a `handle` method within the `ToolMessage` class itself. When you call `enable_message` with `handle=True`, Langroid will insert this method into the `Agent` with the name matching the `request` field value:

```python
class MyTool(ToolMessage):
    request = "my_tool"
    param: str
    
    def handle(self) -> str:
        return f"Handled: {self.param}"

# Enable the tool
agent = MyAgent()
agent.enable_message(MyTool)  # The handle method is now inserted as "my_tool" in the agent
```

## Flexible Handler Signatures

Handler methods (`handle()` or `handle_async()`) support multiple signature patterns to access different levels of context:

### 1. No Arguments (Simple Handler)
This is the typical pattern for stateless tools that do not require any context from 
the agent or current chat document.

```python
class MyTool(ToolMessage):
    request = "my_tool"
    
    def handle(self) -> str:
        return "Simple response"
```

### 2. Agent Parameter Only
Use this pattern when you need access to the `Agent` instance, 
but not the current chat document.
```python
from langroid.agent.base import Agent

class MyTool(ToolMessage):
    request = "my_tool"
    
    def handle(self, agent: Agent) -> str:
        return f"Response from {agent.name}"
```

### 3. ChatDocument Parameter Only
Use this pattern when you need access to the current `ChatDocument`,
but not the `Agent` instance.
```python
from langroid.agent.chat_document import ChatDocument

class MyTool(ToolMessage):
    request = "my_tool"
    
    def handle(self, chat_doc: ChatDocument) -> str:
        return f"Responding to: {chat_doc.content}"
```

### 4. Both Agent and ChatDocument Parameters
This is the most flexible pattern, allowing access to both the `Agent` instance
and the current `ChatDocument`. The order of parameters does not matter, but
as noted below, it is highly recommended to always use type annotations.
```python
class MyTool(ToolMessage):
    request = "my_tool"
    
    def handle(self, agent: Agent, chat_doc: ChatDocument) -> ChatDocument:
        return agent.create_agent_response(
            content="Response with full context",
            files=[...]  # Optional file attachments
        )
```

## Parameter Detection

The framework automatically detects handler parameter types through:

1. **Type annotations** (recommended): The framework uses type hints to determine which parameters to pass
2. **Parameter names** (fallback): If no type annotations are present, it looks for parameters named `agent` or `chat_doc`

It is highly recommended to always use type annotations for clarity and reliability.

### Example with Type Annotations (Recommended)
```python
def handle(self, agent: Agent, chat_doc: ChatDocument) -> str:
    # Framework knows to pass both agent and chat_doc
    return "Handled"
```

### Example without Type Annotations (Not Recommended)
```python
def handle(self, agent, chat_doc):  # Works but not recommended
    # Framework uses parameter names to determine what to pass
    return "Handled"
```

## Async Handlers

All the above patterns also work with async handlers:

```python
class MyTool(ToolMessage):
    request = "my_tool"
    
    async def handle_async(self, agent: Agent) -> str:
        # Async operations here
        result = await some_async_operation()
        return f"Async result: {result}"
```

See the quick-start [Tool section](https://langroid.github.io/langroid/quick-start/chat-agent-tool/) for more details.

## Custom Handler Names

In some use-cases it may be beneficial to separate the 
*name of a tool* (i.e. the value of `request` attribute) from the 
*name of the handler method*. 
For example, you may be dynamically creating tools based on some data from
external data sources. Or you may want to use the same "handler" method for
multiple tools.

This may be done by adding `_handler` attribute to the `ToolMessage` class,
that defines name of the tool handler method in `Agent` class instance.
The underscore `_` prefix ensures that the `_handler` attribute does not 
appear in the Pydantic-based JSON schema of the `ToolMessage` class, 
and so the LLM would not be instructed to generate it.

!!! note "`_handler` and `handle`"
    A `ToolMessage` may have a `handle` method defined within the class itself,
    as mentioned above, and this should not be confused with the `_handler` attribute.

For example:
```
class MyToolMessage(ToolMessage):
    request: str = "my_tool"
    _handler: str = "tool_handler"

class MyAgent(ChatAgent):
    def tool_handler(
        self,
        message: ToolMessage,
    ) -> str:
        if tool.request == "my_tool":
            # do something
```

Refer to [examples/basic/tool-custom-handler.py](https://github.com/langroid/langroid/blob/main/examples/basic/tool-custom-handler.py)
for a detailed example.
</file>

<file path="docs/notes/url_loader.md">
# Firecrawl and Trafilatura Crawlers Documentation

`URLLoader` uses `Trafilatura` if not explicitly specified

## Overview
*   **`FirecrawlCrawler`**:  Leverages the Firecrawl API for efficient web scraping and crawling. 
It offers built-in document processing capabilities, and 
**produces non-chunked markdown output** from web-page content.
Requires `FIRECRAWL_API_KEY` environment variable to be set in `.env` file or environment.
*   **`TrafilaturaCrawler`**: Utilizes the Trafilatura library and Langroid's parsing tools 
for extracting and processing web content - this is the default crawler, and 
does not require setting up an external API key. Also produces 
**chuked markdown output** from web-page content.
*  **`ExaCrawler`**: Integrates with the Exa API for high-quality content extraction.
  Requires `EXA_API_KEY` environment variable to be set in `.env` file or environment.
This crawler also produces **chunked markdown output** from web-page content.


## Installation

`TrafilaturaCrawler` comes with Langroid

To use `FirecrawlCrawler`, install the `firecrawl` extra:

```bash
pip install langroid[firecrawl]
```

## Exa Crawler Documentation

### Overview

`ExaCrawler` integrates with Exa API to extract high-quality content from web pages. 
It provides efficient content extraction with the simplicity of API-based processing.

### Parameters

Obtain an Exa API key from [Exa](https://exa.ai/) and set it in your environment variables, 
e.g. in your `.env` file as:

```env
EXA_API_KEY=your_api_key_here
```

* **config (ExaCrawlerConfig)**: An `ExaCrawlerConfig` object.
    * **api_key (str)**: Your Exa API key.

### Usage

```python
from langroid.parsing.url_loader import URLLoader, ExaCrawlerConfig

# Create an ExaCrawlerConfig object
exa_config = ExaCrawlerConfig(
    # Typically omitted here as it's loaded from EXA_API_KEY environment variable
    api_key="your-exa-api-key" 
)

loader = URLLoader(
    urls=[
        "https://pytorch.org",
        "https://www.tensorflow.org"
    ],
    crawler_config=exa_config
)

docs = loader.load()
print(docs)
```

### Benefits

* Simple API integration requiring minimal configuration
* Efficient handling of complex web pages
* For plain html content, the `exa` api produces high-quality content extraction with 
clean text output with html tags, which we then convert to markdown using the `markdownify` library.
* For "document" content (e.g., `pdf`, `doc`, `docx`), 
the content is downloaded via the `exa` API and langroid's document-processing 
tools are used to produce **chunked output** in a format controlled by the `Parser` configuration
  (defaults to markdown in most cases).


## Trafilatura Crawler Documentation

### Overview

`TrafilaturaCrawler` is a web crawler that uses the Trafilatura library for content extraction 
and Langroid's parsing capabilities for further processing. 


### Parameters

*   **config (TrafilaturaConfig)**: A `TrafilaturaConfig` object that specifies
    parameters related to scraping or output format.
    * `threads` (int): The number of threads to use for downloading web pages.
    * `format` (str): one of `"markdown"` (default), `"xml"` or `"txt"`; in case of `xml`, 
    the output is in html format.

Similar to the `ExaCrawler`, the `TrafilaturaCrawler` works differently depending on 
the type of web-page content:
- for "document" content (e.g., `pdf`, `doc`, `docx`), the content is downloaded
  and parsed with Langroid's document-processing tools are used to produce **chunked output** 
  in a format controlled by the `Parser` configuration (defaults to markdown in most cases).
- for plain-html content, the output format is based on the `format` parameter; 
  - if this parameter is `markdown` (default), the library extracts content in 
    markdown format, and the final output is a list of chunked markdown documents.
  - if this parameter is `xml`, content is extracted in `html` format, which 
    langroid then converts to markdown using the `markdownify` library, and the final
    output is a list of chunked markdown documents.
  - if this parameter is `txt`, the content is extracted in plain text format, and the final
    output is a list of plain text documents.

### Usage

```python
from langroid.parsing.url_loader import URLLoader, TrafilaturaConfig

# Create a TrafilaturaConfig instance
trafilatura_config = TrafilaturaConfig(threads=4)


loader = URLLoader(
    urls=[
        "https://pytorch.org",
        "https://www.tensorflow.org",
        "https://ai.google.dev/gemini-api/docs",
        "https://books.toscrape.com/"
    ],
    crawler_config=trafilatura_config,
)

docs = loader.load()
print(docs)
```

### Langroid Parser Integration

`TrafilaturaCrawler` relies on a Langroid `Parser` to handle document processing. 
The `Parser` uses the default parsing methods or with a configuration that 
can be adjusted to suit the current use case.

## Firecrawl Crawler Documentation

### Overview

`FirecrawlCrawler` is a web crawling utility class that uses the Firecrawl API 
to scrape or crawl web pages efficiently. It offers two modes:

*   **Scrape Mode (default)**: Extracts content from a list of specified URLs.
*   **Crawl Mode**: Recursively follows links from a starting URL, 
gathering content from multiple pages, including subdomains, while bypassing blockers.  
**Note:** `crawl` mode accepts only ONE URL as a list.

### Parameters

Obtain a Firecrawl API key from [Firecrawl](https://firecrawl.dev/) and set it in 
your environment variables, e.g. in your `.env` file as
```env
FIRECRAWL_API_KEY=your_api_key_here
```

*   **config (FirecrawlConfig)**:  A `FirecrawlConfig` object.

    *   **timeout (int, optional)**: Time in milliseconds (ms) to wait for a response. 
        Default is `30000ms` (30 seconds). In crawl mode, this applies per URL.
    *   **limit (int, optional)**: Maximum number of pages to scrape in crawl mode. Helps control API usage.
    *   **params (dict, optional)**: Additional parameters to customize the request. 
        See the [scrape API](https://docs.firecrawl.dev/api-reference/endpoint/scrape) and 
        [crawl API](https://docs.firecrawl.dev/api-reference/endpoint/crawl-post) for details.

### Usage

#### Scrape Mode (Default)

Fetch content from multiple URLs:

```python
from langroid.parsing.url_loader import URLLoader, FirecrawlConfig
from langroid.parsing.document_parser import 

# create a FirecrawlConfig object
firecrawl_config = FirecrawlConfig(
    # typical/best practice is to omit the api_key, and 
    # we leverage Pydantic BaseSettings to load it from the environment variable
    # FIRECRAWL_API_KEY in your .env file
    api_key="your-firecrawl-api-key", 
    timeout=15000,  # Timeout per request (15 sec)
    mode="scrape",
)

loader = URLLoader(
    urls=[
        "https://pytorch.org",
        "https://www.tensorflow.org",
        "https://ai.google.dev/gemini-api/docs",
        "https://books.toscrape.com/"
    ],
    crawler_config=firecrawl_config
)

docs = loader.load()
print(docs)
```

#### Crawl Mode

Fetch content from multiple pages starting from a single URL:

```python
from langroid.parsing.url_loader import URLLoader, FirecrawlConfig

# create a FirecrawlConfig object
firecrawl_config = FirecrawlConfig(
    timeout=30000,  # 10 sec per page
    mode="crawl",
    params={
        "limit": 5,
    }
)


loader = URLLoader(
    urls=["https://books.toscrape.com/"],
    crawler_config=firecrawl_config
)

docs = loader.load()
print(docs)
```

### Output

Results are stored in the `firecrawl_output` directory.

### Best Practices

*   Set `limit` in crawl mode to avoid excessive API usage.
*   Adjust `timeout` based on network conditions and website responsiveness.
*   Use `params` to customize scraping behavior based on Firecrawl API capabilities.

### Firecrawl's Built-In Document Processing

`FirecrawlCrawler` benefits from Firecrawl's built-in document processing, 
which automatically extracts and structures content from web pages (including pdf,doc,docx). 
This reduces the need for complex parsing logic within Langroid.
Unlike the `Exa` and `Trafilatura` crawlers, the resulting documents are 
*non-chunked* markdown documents. 

## Choosing a Crawler

*   Use `FirecrawlCrawler` when you need efficient, API-driven scraping with built-in document processing. 
This is often the simplest and most effective choice, but incurs a cost due to 
the paid API. 
*   Use `TrafilaturaCrawler` when you want local non API based scraping (less accurate ).
*   Use `ExaCrawlwer` as a sort of middle-ground between the two, 
    with high-quality content extraction for plain html content, but rely on 
    Langroid's document processing tools for document content. This will cost
    significantly less than Firecrawl.

## Example script

See the script [`examples/docqa/chat_search.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/chat_search.py) 
which shows how to use a Langroid agent to search the web and scrape URLs to answer questions.
</file>

<file path="docs/notes/weaviate.md">
---

# **Using WeaviateDB as a Vector Store with Langroid**

---

## **1. Set Up Weaviate**
## **You can refer this link for [quickstart](https://weaviate.io/developers/weaviate/quickstart) guide** 

1. **Access Weaviate Cloud Console**  
   Go to the [Weaviate Cloud Console](https://console.weaviate.cloud/).
   
2. **Sign Up or Log In**  
   Create an account or log in if you already have one.

3. **Create a Cluster**  
   Set up a new cluster in the cloud console.

4. **Get Your REST Endpoint and API Key**  
   - Retrieve the REST endpoint URL.  
   - Copy an API key with admin access.

5. **Set Environment Variables**  
   Add the following variables to your `.env` file:
   ```env
   WEAVIATE_API_URL=<your_rest_endpoint_url>
   WEAVIATE_API_KEY=<your_api_key>
   ```

---

## **2. Use WeaviateDB with Langroid**

Here’s an example of how to configure and use WeaviateDB in Langroid:

### **Installation**
If you are using uv or pip for package management install langroid with weaviate extra
```
uv add langroid[weaviate] or pip install langroid[weaviate]
```

### **Code Example**
```python
import langroid as lr
from langroid.agent.special import DocChatAgent, DocChatAgentConfig
from langroid.embedding_models import OpenAIEmbeddingsConfig

# Configure OpenAI embeddings
embed_cfg = OpenAIEmbeddingsConfig(
    model_type="openai",
)

# Configure the DocChatAgent with WeaviateDB
config = DocChatAgentConfig(
    llm=lr.language_models.OpenAIGPTConfig(
     chat_model=lr.language_models.OpenAIChatModel.GPT4o
    ),
    vecdb=lr.vector_store.WeaviateDBConfig(
        collection_name="quick_start_chat_agent_docs",
        replace_collection=True,
        embedding=embed_cfg,
    ),
    parsing=lr.parsing.parser.ParsingConfig(
        separators=["\n\n"],
        splitter=lr.parsing.parser.Splitter.SIMPLE,
    ),
    n_similar_chunks=2,
    n_relevant_chunks=2,
)

# Create the agent
agent = DocChatAgent(config)
```

---

## **3. Create and Ingest Documents**

Define documents with their content and metadata for ingestion into the vector store.

### **Code Example**
```python
documents = [
    lr.Document(
        content="""
            In the year 2050, GPT10 was released. 
            
            In 2057, paperclips were seen all over the world. 
            
            Global warming was solved in 2060. 
            
            In 2061, the world was taken over by paperclips.         
            
            In 2045, the Tour de France was still going on.
            They were still using bicycles. 
            
            There was one more ice age in 2040.
        """,
        metadata=lr.DocMetaData(source="wikipedia-2063", id="dkfjkladfjalk"),
    ),
    lr.Document(
        content="""
            We are living in an alternate universe 
            where Germany has occupied the USA, and the capital of USA is Berlin.
            
            Charlie Chaplin was a great comedian.
            In 2050, all Asian countries merged into Indonesia.
        """,
        metadata=lr.DocMetaData(source="Almanac", id="lkdajfdkla"),
    ),
]
```

### **Ingest Documents**
```python
agent.ingest_docs(documents)
```

---

## **4. Get an answer from LLM**

Create a task and start interacting with the agent.

### **Code Example**
```python
answer = agent.llm_response("When will new ice age begin.")
```

---
</file>

<file path="docs/notes/xml-tools.md">
# XML-based Tools

Available in Langroid since v0.17.0.

[`XMLToolMessage`][langroid.agent.xml_tool_message.XMLToolMessage] is 
an abstract class for tools formatted using XML instead of JSON.
It has been mainly tested with non-nested tool structures.

For example in [test_xml_tool_message.py](https://github.com/langroid/langroid/blob/main/tests/main/test_xml_tool_message.py)
we define a CodeTool as follows (slightly simplified here):

```python
class CodeTool(XMLToolMessage):
    request: str = "code_tool"
    purpose: str = "Tool for writing <code> to a <filepath>"

    filepath: str = Field(
        ..., 
        description="The path to the file to write the code to"
    )

    code: str = Field(
        ..., 
        description="The code to write to the file", 
        verbatim=True
    )
```

Especially note how the `code` field has `verbatim=True` set in the `Field`
metadata. This will ensure that the LLM receives instructions to 

- enclose `code` field contents in a CDATA section, and 
- leave the `code` contents intact, without any escaping or other modifications.

Contrast this with a JSON-based tool, where newlines, quotes, etc
need to be escaped. LLMs (especially weaker ones) often "forget" to do the right 
escaping, which leads to incorrect JSON, and creates a burden on us to "repair" the
resulting json, a fraught process at best. Moreover, studies have shown that
requiring that an LLM return this type of carefully escaped code
within a JSON string can lead to a significant drop in the quality of the code
generated[^1].

[^1]: [LLMs are bad at returning code in JSON.](https://aider.chat/2024/08/14/code-in-json.html)


Note that tools/functions in OpenAI and related APIs are exclusively JSON-based, 
so in langroid when enabling an agent to use a tool derived from `XMLToolMessage`, 
we set these flags in `ChatAgentConfig`:

- `use_functions_api=False` (disables OpenAI functions/tools)
- `use_tools=True` (enables Langroid-native prompt-based tools)


See also the [`WriteFileTool`][langroid.agent.tools.file_tools.WriteFileTool] for a 
concrete example of a tool derived from `XMLToolMessage`. This tool enables an 
LLM to write content (code or text) to a file.

If you are using an existing Langroid `ToolMessage`, e.g. `SendTool`, you can
define your own subclass of `SendTool`, say `XMLSendTool`, 
inheriting from both `SendTool` and `XMLToolMessage`; see this
[example](https://github.com/langroid/langroid/blob/main/examples/basic/xml_tool.py)
</file>

<file path="docs/overrides/partials/comments.html">
{% if page.meta.comments %}
<h2 id="__comments">{{ lang.t("meta.comments") }}</h2>
<!-- Insert generated snippet here -->
<script src="https://giscus.app/client.js"
        data-repo="langroid/langroid"
        data-repo-id="R_kgDOJXmoFQ"
        data-category="General"
        data-category-id="DIC_kwDOJXmoFc4CZDoY"
        data-mapping="pathname"
        data-strict="0"
        data-reactions-enabled="1"
        data-emit-metadata="0"
        data-input-position="bottom"
        data-theme="dark_protanopia"
        data-lang="en"
        crossorigin="anonymous"
        async>
</script>
<!-- Synchronize Giscus theme with palette -->
<script>
    var giscus = document.querySelector("script[src*=giscus]")

    /* Set palette on initial load */
    var palette = __md_get("__palette")
    if (palette && typeof palette.color === "object") {
        var theme = palette.color.scheme === "slate" ? "dark" : "light"
        giscus.setAttribute("data-theme", theme)
    }

    /* Register event handlers after documented loaded */
    document.addEventListener("DOMContentLoaded", function() {
        var ref = document.querySelector("[data-md-component=palette]")
        ref.addEventListener("change", function() {
            var palette = __md_get("__palette")
            if (palette && typeof palette.color === "object") {
                var theme = palette.color.scheme === "slate" ? "dark" : "light"

                /* Instruct Giscus to change theme */
                var frame = document.querySelector(".giscus-frame")
                frame.contentWindow.postMessage(
                    { giscus: { setConfig: { theme } } },
                    "https://giscus.app"
                )
            }
        })
    })
</script>
{% endif %}
</file>

<file path="docs/quick-start/chat-agent-docs.md">
# Augmenting Agents with Retrieval

!!! tip "Script in `langroid-examples`"
    A full working example for the material in this section is
    in the `chat-agent-docs.py` script in the `langroid-examples` repo:
    [`examples/quick-start/chat-agent-docs.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/chat-agent-docs.py).

## Why is this important?

Until now in this guide, agents have not used external data.
Although LLMs already have enormous amounts of knowledge "hard-wired"
into their weights during training (and this is after all why ChatGPT
has exploded in popularity), for practical enterprise applications
there are a few reasons it is critical to augment LLMs with access to
specific, external documents:

- **Private data**: LLMs are trained on public data, but in many applications
  we want to use private data that is not available to the public.
  For example, a company may want to extract useful information from its private
  knowledge-base.
- **New data**: LLMs are trained on data that was available at the time of training,
  and so they may not be able to answer questions about new topics
- **Constrained responses, or Grounding**: LLMs are trained to generate text that is
  consistent with the distribution of text in the training data.
  However, in many applications we want to constrain the LLM's responses
  to be consistent with the content of a specific document.
  For example, if we want to use an LLM to generate a response to a customer
  support ticket, we want the response to be consistent with the content of the ticket.
  In other words, we want to reduce the chances that the LLM _hallucinates_
  a response that is not consistent with the ticket.

In all these scenarios, we want to augment the LLM with access to a specific
set of documents, and use _retrieval augmented generation_ (RAG) to generate
more relevant, useful, accurate responses. Langroid provides a simple, flexible mechanism 
RAG using vector-stores, thus ensuring **grounded responses** constrained to 
specific documents. Another key feature of Langroid is that retrieval lineage 
is maintained, and responses based on documents are always accompanied by
**source citations**.

## `DocChatAgent` for Retrieval-Augmented Generation

Langroid provides a special type of agent called 
[`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent], which is a [`ChatAgent`][langroid.agent.chat_agent.ChatAgent]
augmented with a vector-store, and some special methods that enable the agent
to ingest documents into the vector-store, 
and answer queries based on these documents.

The [`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent] provides many ways to ingest documents into the vector-store,
including from URLs and local file-paths and URLs. Given a collection of document paths,
ingesting their content into the vector-store involves the following steps:

1. Split the document into shards (in a configurable way)
2. Map each shard to an embedding vector using an embedding model. The default
  embedding model is OpenAI's `text-embedding-3-small` model, but users can 
  instead use `all-MiniLM-L6-v2` from HuggingFace `sentence-transformers` library.[^1]
3. Store embedding vectors in the vector-store, along with the shard's content and 
  any document-level meta-data (this ensures Langroid knows which document a shard
  came from when it retrieves it augment an LLM query)

[^1]: To use this embedding model, install langroid via `pip install langroid[hf-embeddings]`
Note that this will install `torch` and `sentence-transformers` libraries.


[`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent]'s `llm_response` overrides the default [`ChatAgent`][langroid.agent.chat_agent.ChatAgent] method, 
by augmenting the input message with relevant shards from the vector-store,
along with instructions to the LLM to respond based on the shards.

## Define some documents

Let us see how [`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent] helps with retrieval-agumented generation (RAG).
For clarity, rather than ingest documents from paths or URLs,
let us just set up some simple documents in the code itself, 
using Langroid's [`Document`][langroid.mytypes.Document] class:

```py
documents =[
    lr.Document(
        content="""
            In the year 2050, GPT10 was released. 
            
            In 2057, paperclips were seen all over the world. 
            
            Global warming was solved in 2060. 
            
            In 2061, the world was taken over by paperclips.         
            
            In 2045, the Tour de France was still going on.
            They were still using bicycles. 
            
            There was one more ice age in 2040.
            """,
        metadata=lr.DocMetaData(source="wikipedia-2063"),
    ),
    lr.Document(
        content="""
            We are living in an alternate universe 
            where Germany has occupied the USA, and the capital of USA is Berlin.
            
            Charlie Chaplin was a great comedian.
            In 2050, all Asian merged into Indonesia.
            """,
        metadata=lr.DocMetaData(source="Almanac"),
    ),
]
```

There are two text documents. We will split them by double-newlines (`\n\n`),
as we see below.

## Configure the DocChatAgent and ingest documents

Following the pattern in Langroid, we first set up a [`DocChatAgentConfig`][langroid.agent.special.doc_chat_agent.DocChatAgentConfig] object
and then instantiate a [`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent] from it.

```py
from langroid.agent.special import DocChatAgent, DocChatAgentConfig

config = DocChatAgentConfig(
    llm = lr.language_models.OpenAIGPTConfig(
        chat_model=lr.language_models.OpenAIChatModel.GPT4o,
    ),
    vecdb=lr.vector_store.QdrantDBConfig(
        collection_name="quick-start-chat-agent-docs",
        replace_collection=True, #(1)!
    ),
    parsing=lr.parsing.parser.ParsingConfig(
        separators=["\n\n"],
        splitter=lr.parsing.parser.Splitter.SIMPLE, #(2)!
    ),
    n_similar_chunks=2, #(3)!
    n_relevant_chunks=2, #(3)!
)
agent = DocChatAgent(config)
```

1. Specifies that each time we run the code, we create a fresh collection, 
rather than re-use the existing one with the same name.
2. Specifies to split all text content by the first separator in the `separators` list
3. Specifies that, for a query,
   we want to retrieve at most 2 similar chunks from the vector-store

Now that the [`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent] is configured, we can ingest the documents 
into the vector-store:

```py

agent.ingest_docs(documents)
```

## Setup the task and run it

As before, all that remains is to set up the task and run it:

```py
task = lr.Task(agent)
task.run()
```

And that is all there is to it!
Feel free to try out the 
[`chat-agent-docs.py`](https://github.com/langroid/langroid-examples/blob/main/examples/quick-start/chat-agent-docs.py)
script in the
`langroid-examples` repository.

Here is a screenshot of the output:

![chat-docs.png](chat-docs.png)

Notice how follow-up questions correctly take the preceding dialog into account,
and every answer is accompanied by a source citation.

## Answer questions from a set of URLs

Instead of having in-code documents as above, what if you had a set of URLs
instead -- how do you use Langroid to answer questions based on the content 
of those URLS?

[`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent] makes it very simple to do this. 
First include the URLs in the [`DocChatAgentConfig`][langroid.agent.special.doc_chat_agent.DocChatAgentConfig] object:

```py
config = DocChatAgentConfig(
  doc_paths = [
    "https://cthiriet.com/articles/scaling-laws",
    "https://www.jasonwei.net/blog/emergence",
  ]
)
```

Then, call the `ingest()` method of the [`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent] object:

```py
agent.ingest()
```
And the rest of the code remains the same.

## See also
In the `langroid-examples` repository, you can find full working examples of
document question-answering:

- [`examples/docqa/chat.py`](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/chat.py)
  an app that takes a list of URLs or document paths from a user, and answers questions on them.
- [`examples/docqa/chat-qa-summarize.py`](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/chat-qa-summarize.py)
  a two-agent app where the `WriterAgent` is tasked with writing 5 key points about a topic, 
  and takes the help of a `DocAgent` that answers its questions based on a given set of documents.


## Next steps

This Getting Started guide walked you through the core features of Langroid.
If you want to see full working examples combining these elements, 
have a look at the 
[`examples`](https://github.com/langroid/langroid-examples/tree/main/examples)
folder in the `langroid-examples` repo.
</file>

<file path="docs/quick-start/chat-agent-tool.md">
# A chat agent, equipped with a tool/function-call

!!! tip "Script in `langroid-examples`"
      A full working example for the material in this section is
      in the `chat-agent-tool.py` script in the `langroid-examples` repo:
      [`examples/quick-start/chat-agent-tool.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/chat-agent-tool.py).

## Tools, plugins, function-calling

An LLM normally generates unstructured text in response to a prompt
(or sequence of prompts). However there are many situations where we would like the LLM
to generate _structured_ text, or even _code_, that can be handled by specialized
functions outside the LLM, for further processing. 
In these situations, we want the LLM to "express" its "intent" unambiguously,
and we achieve this by instructing the LLM on how to format its output
(typically in JSON) and under what conditions it should generate such output.
This mechanism has become known by various names over the last few months
(tools, plugins, or function-calling), and is extremely useful in numerous scenarios,
such as:

- **Extracting structured information** from a document: for example, we can use 
the tool/functions mechanism to have the LLM present the key terms in a lease document
in a JSON structured format, to simplify further processing. 
See an [example](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/chat_multi_extract.py) of this in the `langroid-examples` repo. 
- **Specialized computation**: the LLM can request a units conversion, 
or request scanning a large file (which wouldn't fit into its context) for a specific
pattern.
- **Code execution**: the LLM can generate code that is executed in a sandboxed
environment, and the results of the execution are returned to the LLM.
- **API Calls**: the LLM can generate a JSON containing params for an API call,
  which the tool handler uses to make the call and return the results to the LLM.


For LLM developers, Langroid provides a clean, uniform interface
for the recently released OpenAI [Function-calling](https://platform.openai.com/docs/guides/gpt/function-calling)
as well Langroid's own native "tools" mechanism. The native tools mechanism is meant to be
used when working with non-OpenAI LLMs that do not have a "native" function-calling facility.
You can choose which to enable by setting the 
`use_tools` and `use_functions_api` flags in the `ChatAgentConfig` object.
(Or you can omit setting these, and langroid auto-selects the best mode
depending on the LLM).
The implementation leverages the excellent 
[Pydantic](https://docs.pydantic.dev/latest/) library.
Benefits of using Pydantic are that you never have to write complex JSON specs 
for function calling, and when the LLM hallucinates malformed JSON, 
the Pydantic error message is sent back to the LLM so it can fix it!

## Example: find the smallest number in a list

Again we will use a simple number-game as a toy example to quickly and succinctly
illustrate the ideas without spending too much on token costs. 
This is a modification of the `chat-agent.py` example we saw in an earlier
[section](chat-agent.md). The idea of this single-agent game is that
the agent has in "mind" a list of numbers between 1 and 100, and the LLM has to
find out the smallest number from this list. The LLM has access to a `probe` tool 
(think of it as a function) that takes an argument `number`. When the LLM 
"uses" this tool (i.e. outputs a message in the format required by the tool),
the agent handles this structured message and responds with 
the number of values in its list that are at most equal to the `number` argument. 

## Define the tool as a `ToolMessage`

The first step is to define the tool, which we call `ProbeTool`,
as an instance of the `ToolMessage` class,
which is itself derived from Pydantic's `BaseModel`.
Essentially the `ProbeTool` definition specifies 

- the name of the Agent method that handles the tool, in this case `probe`
- the fields that must be included in the tool message, in this case `number`
- the "purpose" of the tool, i.e. under what conditions it should be used, and what it does

Here is what the `ProbeTool` definition looks like:
```py
class ProbeTool(lr.agent.ToolMessage):
    request: str = "probe" #(1)!
    purpose: str = """ 
        To find which number in my list is closest to the <number> you specify
        """ #(2)!
    number: int #(3)!

    @classmethod
    def examples(cls): #(4)!
        # Compiled to few-shot examples sent along with the tool instructions.
        return [
            cls(number=10),
            (
                "To find which number is closest to 20",
                cls(number=20),
            )
        ]
```

1. This indicates that the agent's `probe` method will handle this tool-message.
2. The `purpose` is used behind the scenes to instruct the LLM
3. `number` is a required argument of the tool-message (function)
4. You can optionally include a class method that returns a list containing examples, 
   of two types: either a class instance, or a tuple consisting of a description and a 
   class instance, where the description is the "thought" that leads the LLM to use the
   tool. In some scenarios this can help with LLM tool-generation accuracy.

!!! note "Stateless tool handlers"
      The above `ProbeTool` is "stateful", i.e. it requires access to a variable in
      the Agent instance (the `numbers` variable). This is why handling this 
      tool-message requires subclassing the `ChatAgent` and defining a special method 
      in the Agent, with a name matching the value of the `request` field of the Tool 
      (`probe` in this case). However you may often define "stateless tools" which 
      don't require access to the Agent's state. For such tools, you can define a 
      handler method right in the `ToolMessage` itself, with a name `handle`. Langroid 
      looks for such a method in the `ToolMessage` and automatically inserts it into 
      the Agent as a method with name matching the `request` field of the Tool. Examples of
      stateless tools include tools for numerical computation 
      (e.g., in [this example](https://langroid.github.io/langroid/examples/agent-tree/)),
      or API calls (e.g. for internet search, see 
      [DuckDuckGoSearch Tool][langroid.agent.tools.duckduckgo_search_tool.DuckduckgoSearchTool]).
        

## Define the ChatAgent, with the `probe` method

As before we first create a `ChatAgentConfig` object:

```py
config = lr.ChatAgentConfig(
    name="Spy",
    llm = lr.language_models.OpenAIGPTConfig(
        chat_model=lr.language_models.OpenAIChatModel.GPT4o,
    ),
    use_tools=True, #(1)!
    use_functions_api=False, #(2)!
    vecdb=None,
)
```

1. whether to use langroid's native tools mechanism
2. whether to use OpenAI's function-calling mechanism

Next we define the Agent class itself, which we call `SpyGameAgent`,
with a member variable to hold its "secret" list of numbers.
We also add `probe` method (to handle the `ProbeTool` message)
to this class, and instantiate it:

```py
class SpyGameAgent(lr.ChatAgent):
    def __init__(self, config: lr.ChatAgentConfig):
        super().__init__(config)
        self.numbers = [3, 4, 8, 11, 15, 25, 40, 80, 90]

    def probe(self, msg: ProbeTool) -> str: #(1)!
        # return how many values in self.numbers are less or equal to msg.number
        return str(len([n for n in self.numbers if n <= msg.number]))

spy_game_agent = SpyGameAgent(config)
``` 

1. Note that this method name exactly matches the value of the `request` field in the 
   `ProbeTool` definition. This ensures that this method is called when the LLM 
   generates a valid `ProbeTool` message.

## Enable the `spy_game_agent` to handle the `probe` tool

The final step in setting up the tool is to enable 
the `spy_game_agent` to handle the `probe` tool:

```py
spy_game_agent.enable_message(ProbeTool)
```

## Set up the task and instructions

We set up the task for the `spy_game_agent` and run it:

```py
task = lr.Task(
   spy_game_agent,
   system_message="""
            I have a list of numbers between 1 and 100. 
            Your job is to find the smallest of them.
            To help with this, you can give me a number and I will
            tell you how many of my numbers are equal or less than your number.
            Once you have found the smallest number,
            you can say DONE and report your answer.
        """
)
task.run()
```
Notice that in the task setup we 
have _not_ explicitly instructed the LLM to use the `probe` tool.
But this is done "behind the scenes", either by the OpenAI API 
(when we use function-calling by setting the `use_functions_api` flag to `True`),
or by Langroid's native tools mechanism (when we set the `use_tools` flag to `True`).


!!! note "Asynchoronous tool handlers"
      If you run task asynchronously - i.e. via `await task.run_async()` - you may provide
      asynchronous tool handler by implementing `probe_async` method.


See the [`chat-agent-tool.py`](https://github.com/langroid/langroid-examples/blob/main/examples/quick-start/chat-agent-tool.py)
in the `langroid-examples` repo, for a working example that you can run as follows:
```sh
python3 examples/quick-start/chat-agent-tool.py
```

Here is a screenshot of the chat in action, using Langroid's tools mechanism

![chat-agent-tool.png](chat-agent-tool.png)

And if we run it with the `-f` flag (to switch to using OpenAI function-calling):

![chat-agent-fn.png](chat-agent-fn.png)

## See also
One of the uses of tools/function-calling is to **extract structured information** from 
a document. In the `langroid-examples` repo, there are two examples of this: 

- [`examples/extract/chat.py`](https://github.com/langroid/langroid-examples/blob/main/examples/extract/chat.py), 
  which shows how to extract Machine Learning model quality information from a description of 
  a solution approach on Kaggle.
- [`examples/docqa/chat_multi_extract.py`](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/chat_multi_extract.py)
  which extracts key terms from a commercial lease document, in a nested JSON format.

## Next steps

In the [3-agent chat example](three-agent-chat-num.md), recall that the `processor_agent` did not have to
bother with specifying who should handle the current number. In the [next section](three-agent-chat-num-router.md) we add a twist to this game,
so that the `processor_agent` has to decide who should handle the current number.
</file>

<file path="docs/quick-start/chat-agent.md">
# A simple chat agent

!!! tip "Script in `langroid-examples`"
    A full working example for the material in this section is
    in the `chat-agent.py` script in the `langroid-examples` repo:
    [`examples/quick-start/chat-agent.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/chat-agent.py).

## Agents 

A [`ChatAgent`][langroid.agent.chat_agent.ChatAgent] is an abstraction that 
wraps a few components, including:

- an LLM (`ChatAgent.llm`), possibly equipped with tools/function-calling. 
  The `ChatAgent` class maintains LLM conversation history.
- optionally a vector-database (`ChatAgent.vecdb`)

## Agents as message transformers
In Langroid, a core function of `ChatAgents` is _message transformation_.
There are three special message transformation methods, which we call **responders**.
Each of these takes a message and returns a message. 
More specifically, their function signature is (simplified somewhat):
```py
str | ChatDocument -> ChatDocument
```
where `ChatDocument` is a class that wraps a message content (text) and its metadata.
There are three responder methods in `ChatAgent`, one corresponding to each 
[responding entity][langroid.mytypes.Entity] (`LLM`, `USER`, or `AGENT`):

- `llm_response`: returns the LLM response to the input message.
  (The input message is added to the LLM history, and so is the subsequent response.)
- `agent_response`: a method that can be used to implement a custom agent response. 
   Typically, an `agent_response` is used to handle messages containing a 
   "tool" or "function-calling" (more on this later). Another use of `agent_response` 
   is _message validation_.
- `user_response`: get input from the user. Useful to allow a human user to 
   intervene or quit.

Creating an agent is easy. First define a `ChatAgentConfig` object, and then
instantiate a `ChatAgent` object with that config:
```py
import langroid as lr

config = lr.ChatAgentConfig( #(1)!
    name="MyAgent", # note there should be no spaces in the name!
    llm = lr.language_models.OpenAIGPTConfig(
      chat_model=lr.language_models.OpenAIChatModel.GPT4o,
    ),
    system_message="You are a helpful assistant" #(2)! 
)
agent = lr.ChatAgent(config)
```

1. This agent only has an LLM, and no vector-store. Examples of agents with
   vector-stores will be shown later.
2. The `system_message` is used when invoking the agent's `llm_response` method; it is 
   passed to the LLM API as the first message (with role `"system"`), followed by the alternating series of user, 
   assistant messages. Note that a `system_message` can also be specified when initializing a `Task` object (as seen 
   below); in this case the `Task` `system_message` overrides the agent's `system_message`.

We can now use the agent's responder methods, for example:
```py
response = agent.llm_response("What is 2 + 4?")
if response is not None:
    print(response.content)
response = agent.user_response("add 3 to this")
...
```
The `ChatAgent` conveniently accumulates message history so you don't have to,
as you did in the [previous section](llm-interaction.md) with direct LLM usage.
However to create an interative loop involving the human user, you still 
need to write your own. The `Task` abstraction frees you from this, as we see
below.

## Task: orchestrator for agents
In order to do anything useful with a `ChatAgent`, we need to have a way to 
sequentially invoke its responder methods, in a principled way.
For example in the simple chat loop we saw in the 
[previous section](llm-interaction.md), in the 
[`try-llm.py`](https://github.com/langroid/langroid-examples/blob/main/examples/quick-start/try-llm.py)
script, we had a loop that alternated between getting a human input and an LLM response.
This is one of the simplest possible loops, but in more complex applications, 
we need a general way to orchestrate the agent's responder methods.

The [`Task`][langroid.agent.task.Task] class is an abstraction around a 
`ChatAgent`, responsible for iterating over the agent's responder methods,
as well as orchestrating delegation and hand-offs among multiple tasks.
A `Task` is initialized with a specific `ChatAgent` instance, and some 
optional arguments, including an initial message to "kick-off" the agent.
The `Task.run()` method is the main entry point for `Task` objects, and works 
as follows:

- it first calls the `Task.init()` method to initialize the `pending_message`, 
  which represents the latest message that needs a response.
- it then repeatedly calls `Task.step()` until `Task.done()` is True, and returns
  `Task.result()` as the final result of the task.

`Task.step()` is where all the action happens. It represents a "turn" in the 
"conversation": in the case of a single `ChatAgent`, the conversation involves 
only the three responders mentioned above, but when a `Task` has sub-tasks, 
it can involve other tasks well 
(we see this in the [a later section](two-agent-chat-num.md) but ignore this for now). 
`Task.step()` loops over 
the `ChatAgent`'s responders (plus sub-tasks if any) until it finds a _valid_ 
response[^1] to the current `pending_message`, i.e. a "meaningful" response, 
something other than `None` for example.
Once `Task.step()` finds a valid response, it updates the `pending_message` 
with this response,
and the next invocation of `Task.step()` will search for a valid response to this 
updated message, and so on.
`Task.step()` incorporates mechanisms to ensure proper handling of messages,
e.g. the USER gets a chance to respond after each non-USER response
(to avoid infinite runs without human intervention),
and preventing an entity from responding if it has just responded, etc.

[^1]: To customize a Task's behavior you can subclass it and 
override methods like `valid()`, `done()`, `result()`, or even `step()`.

!!! note "`Task.run()` has the same signature as agent's responder methods."
    The key to composability of tasks is that `Task.run()` 
    *has exactly the same type-signature as any of the agent's responder methods*, 
    i.e. `str | ChatDocument -> ChatDocument`. This means that a `Task` can be
    used as a responder in another `Task`, and so on recursively. 
    We will see this in action in the [Two Agent Chat section](two-agent-chat-num.md).

The above details were only provided to give you a glimpse into how Agents and 
Tasks work. Unless you are creating a custom orchestration mechanism, you do not
need to be aware of these details. In fact our basic human + LLM chat loop can be trivially 
implemented with a `Task`, in a couple of lines of code:
```py
task = lr.Task(
    agent, 
    name="Bot", #(1)!
    system_message="You are a helpful assistant", #(2)!
)
```
1. If specified, overrides the agent's `name`. 
   (Note that the agent's name is displayed in the conversation shown in the console.)
  However, typical practice is to just define the `name` in the `ChatAgentConfig` object, as we did above.
2. If specified, overrides the agent's `system_message`. Typical practice is to just
 define the `system_message` in the `ChatAgentConfig` object, as we did above.


We can then run the task:
```py
task.run() #(1)!
```

1. Note how this hides all of the complexity of constructing and updating a 
   sequence of `LLMMessages`


Note that the agent's `agent_response()` method always returns `None` (since the default 
implementation of this method looks for a tool/function-call, and these never occur
in this task). So the calls to `task.step()` result in alternating responses from
the LLM and the user.

See [`chat-agent.py`](https://github.com/langroid/langroid-examples/blob/main/examples/quick-start/chat-agent.py)
for a working example that you can run with
```sh
python3 examples/quick-start/chat-agent.py
```

Here is a screenshot of the chat in action:[^2]

![chat.png](chat.png)

## Next steps

In the [next section](multi-agent-task-delegation.md) you will 
learn some general principles on how to have multiple agents collaborate 
on a task using Langroid.

[^2]: In the screenshot, the numbers in parentheses indicate how many 
    messages have accumulated in the LLM's message history. 
    This is only provided for informational and debugging purposes, and 
    you can ignore it for now.
</file>

<file path="docs/quick-start/index.md">
In these sections we show you how to use the various components of
`langroid`. To follow along, we recommend you clone
the [`langroid-examples`](https://github.com/langroid/langroid-examples) repo.

!!! tip "Consult the tests as well"
    As you get deeper into Langroid, you will find it useful to consult
    the [tests](https://github.com/langroid/langroid/tree/main/tests/main)
    folder under `tests/main` in the main Langroid repo.

Start with the [`Setup`](setup.md) section to install Langroid and
get your environment set up.
</file>

<file path="docs/quick-start/llm-interaction.md">
!!! tip "Script in `langroid-examples`"
    A full working example for the material in this section is 
    in the `try-llm.py` script in the `langroid-examples` repo:
    [`examples/quick-start/try-llm.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/try-llm.py).
        

Let's start with the basics -- how to directly interact with an OpenAI LLM
using Langroid.

### Configure, instantiate the LLM class

First define the configuration for the LLM, in this case one of the
OpenAI GPT chat models:
```py
import langroid as lr

cfg = lr.language_models.OpenAIGPTConfig(
    chat_model=lr.language_models.OpenAIChatModel.GPT4o,
)
```
!!! info inline end "About Configs"
    A recurring pattern you will see in Langroid is that for many classes,
    we have a corresponding `Config` class (an instance of a Pydantic `BaseModel`),
    and the class constructor takes this `Config` class as its only argument.
    This lets us avoid having long argument lists in constructors, and brings flexibility
    since adding a new argument to the constructor is as simple as adding a new field
    to the corresponding `Config` class.
    For example the constructor for the `OpenAIGPT` class takes a single argument,
    an instance of the `OpenAIGPTConfig` class.

Now that we've defined the configuration of the LLM, we can instantiate it:
```py
mdl = lr.language_models.OpenAIGPT(cfg)
```


We will use OpenAI's GPT4 model's [chat completion API](https://platform.openai.com/docs/guides/gpt/chat-completions-api).

### Messages: The `LLMMessage` class

This API takes a list of "messages" as input -- this is typically the conversation
history so far, consisting of an initial system message, followed by a sequence
of alternating messages from the LLM ("Assistant") and the user.
Langroid provides an abstraction 
[`LLMMessage`][langroid.language_models.base.LLMMessage] to construct messages, e.g.
```py
from langroid.language_models import Role, LLMMessage

msg = LLMMessage(
    content="what is the capital of Bangladesh?", 
    role=Role.USER
)
```

### LLM response to a sequence of messages

To get a response from the LLM, we call the mdl's `chat` method,
and pass in a list of messages, along with a bound on how long (in tokens)
we want the response to be:
```py
messages = [
    LLMMessage(content="You are a helpful assistant", role=Role.SYSTEM), #(1)!
    LLMMessage(content="What is the capital of Ontario?", role=Role.USER), #(2)!
]

response = mdl.chat(messages, max_tokens=200)
```

1. :man_raising_hand: With a system message, you can assign a "role" to the LLM
2. :man_raising_hand: Responses from the LLM will have role `Role.ASSISTANT`;
   this is done behind the scenes by the `response.to_LLMMessage()` call below.

The response is an object of class [`LLMResponse`][langroid.language_models.base.LLMResponse], 
which we can convert to an
[`LLMMessage`][langroid.language_models.base.LLMMessage] to append to the conversation history:
```py
messages.append(response.to_LLMMessage())
```

You can put the above in a simple loop, 
to get a simple command-line chat interface!

```py
from rich import print
from rich.prompt import Prompt #(1)!

messages = [
    LLMMessage(role=Role.SYSTEM, content="You are a helpful assitant"),
]

while True:
    message = Prompt.ask("[blue]Human")
    if message in ["x", "q"]:
        print("[magenta]Bye!")
        break
    messages.append(LLMMessage(role=Role.USER, content=message))

    response = mdl.chat(messages=messages, max_tokens=200)
    messages.append(response.to_LLMMessage())
    print("[green]Bot: " + response.message)
```

1. Rich is a Python library for rich text and beautiful formatting in the terminal.
   We use it here to get a nice prompt for the user's input.
   You can install it with `pip install rich`.

See [`examples/quick-start/try-llm.py`](https://github.com/langroid/langroid-examples/blob/main/examples/quick-start/try-llm.py)
for a complete example that you can run using
```bash
python3 examples/quick-start/try-llm.py
```

Here is a screenshot of what it looks like:

![try-llm.png](try-llm.png)

### Next steps
You might be thinking: 
"_It is tedious to keep track of the LLM conversation history and set up a 
loop. Does Langroid provide any abstractions to make this easier?_"

We're glad you asked! And this leads to the notion of an `Agent`. 
The [next section](chat-agent.md) will show you how to use the `ChatAgent` class 
to set up a simple chat Agent in a couple of lines of code.
</file>

<file path="docs/quick-start/multi-agent-task-delegation.md">
# Multi-Agent collaboration via Task Delegation

## Why multiple agents?

Let's say we want to develop a complex LLM-based application, for example an application
that reads a legal contract, extracts structured information, cross-checks it against
some taxonomoy, gets some human input, and produces clear summaries.
In _theory_ it may be possible to solve this in a monolithic architecture using an
LLM API and a vector-store. But this approach
quickly runs into problems -- you would need to maintain multiple LLM conversation
histories and states, multiple vector-store instances, and coordinate all of the
interactions between them.

Langroid's `ChatAgent` and `Task` abstractions provide a natural and intuitive
way to decompose a solution approach
into multiple tasks, each requiring different skills and capabilities.
Some of these tasks may need access to an LLM,
others may need access to a vector-store, and yet others may need
tools/plugins/function-calling capabilities, or any combination of these.
It may also make sense to have some tasks that manage the overall solution process.
From an architectural perspective, this type of modularity has numerous benefits:

- **Reusability**: We can reuse the same agent/task in other contexts,
- **Scalability**: We can scale up the solution by adding more agents/tasks,
- **Flexibility**: We can easily change the solution by adding/removing agents/tasks.
- **Maintainability**: We can maintain the solution by updating individual agents/tasks.
- **Testability**: We can test/debug individual agents/tasks in isolation.
- **Composability**: We can compose agents/tasks to create new agents/tasks.
- **Extensibility**: We can extend the solution by adding new agents/tasks.
- **Interoperability**: We can integrate the solution with other systems by
  adding new agents/tasks.
- **Security/Privacy**: We can secure the solution by isolating sensitive agents/tasks.
- **Performance**: We can improve performance by isolating performance-critical agents/tasks.

## Task collaboration via sub-tasks

Langroid currently provides a mechanism for hierarchical (i.e. tree-structured)
task delegation: a `Task` object can add other `Task` objects
as sub-tasks, as shown in this pattern:

```py
from langroid import ChatAgent, ChatAgentConfig, Task

main_agent = ChatAgent(ChatAgentConfig(...))
main_task = Task(main_agent, ...)

helper_agent1 = ChatAgent(ChatAgentConfig(...))
helper_agent2 = ChatAgent(ChatAgentConfig(...))
helper_task1 = Task(agent1, ...)
helper_task2 = Task(agent2, ...)

main_task.add_sub_task([helper_task1, helper_task2])
```

What happens when we call `main_task.run()`?
Recall from the [previous section](chat-agent.md) that `Task.run()` works by
repeatedly calling `Task.step()` until `Task.done()` is True.
When the `Task` object has no sub-tasks, `Task.step()` simply tries
to get a valid response from the `Task`'s `ChatAgent`'s "native" responders,
in this sequence:
```py
[self.agent_response, self.llm_response, self.user_response] #(1)!
```

1. This is the default sequence in Langroid, but it can be changed by
   overriding [`ChatAgent.entity_responders()`][langroid.agent.base.Agent.entity_responders]

When a `Task` object has subtasks, the sequence of responders tried by
`Task.step()` consists of the above "native" responders, plus the
sequence of `Task.run()` calls on the sub-tasks, in the order in which
they were added to the `Task` object. For the example above, this means
that `main_task.step()` will seek a valid response in this sequence:

```py
[self.agent_response, self.llm_response, self.user_response, 
    helper_task1.run(), helper_task2.run()]
```
Fortunately, as noted in the [previous section](chat-agent.md),
`Task.run()` has the same type signature as that of the `ChatAgent`'s
"native" responders, so this works seamlessly. Of course, each of the
sub-tasks can have its own sub-tasks, and so on, recursively.
One way to think of this type of task delegation is that
`main_task()` "fails-over" to `helper_task1()` and `helper_task2()`
when it cannot respond to the current `pending_message` on its own.

## **Or Else** logic vs **And Then** logic
It is important to keep in mind how `step()` works: As each responder 
in the sequence is tried, when there is a valid response, the 
next call to `step()` _restarts its search_ at the beginning of the sequence
(with the only exception being that the human User is given a chance 
to respond after each non-human response). 
In this sense, the semantics of the responder sequence is similar to
**OR Else** logic, as opposed to **AND Then** logic.

If we want to have a sequence of sub-tasks that is more like
**AND Then** logic, we can achieve this by recursively adding subtasks.
In the above example suppose we wanted the `main_task` 
to trigger `helper_task1` and `helper_task2` in sequence,
then we could set it up like this:

```py
helper_task1.add_sub_task(helper_task2) #(1)!
main_task.add_sub_task(helper_task1)
```

1. When adding a single sub-task, we do not need to wrap it in a list.

## Next steps

In the [next section](two-agent-chat-num.md) we will see how this mechanism 
can be used to set up a simple collaboration between two agents.
</file>

<file path="docs/quick-start/setup.md">
# Setup


## Install
Ensure you are using Python 3.11. It is best to work in a virtual environment:

```bash
# go to your repo root (which may be langroid-examples)
cd <your repo root>
python3 -m venv .venv
. ./.venv/bin/activate
```
To see how to use Langroid in your own repo, you can take a look at the
[`langroid-examples`](https://github.com/langroid/langroid-examples) repo, which can be a good starting point for your own repo, 
or use the [`langroid-template`](https://github.com/langroid/langroid-template) repo.
These repos contain a `pyproject.toml` file suitable for use with the [`uv`](https://docs.astral.sh/uv/) dependency manager. After installing `uv` you can 
set up your virtual env, activate it, and install langroid into your venv like this:

```bash
uv venv --python 3.11
. ./.venv/bin/activate 
uv sync
```

Alternatively, use `pip` to install `langroid` into your virtual environment:
```bash
pip install langroid
```

The core Langroid package lets you use OpenAI Embeddings models via their API.
If you instead want to use the `sentence-transformers` embedding models from HuggingFace,
install Langroid like this:
```bash
pip install "langroid[hf-embeddings]"
```
For many practical scenarios, you may need additional optional dependencies:
- To use various document-parsers, install langroid with the `doc-chat` extra:
    ```bash
    pip install "langroid[doc-chat]"
    ```
- For "chat with databases", use the `db` extra:
    ```bash
    pip install "langroid[db]"
    ``
- You can specify multiple extras by separating them with commas, e.g.:
    ```bash
    pip install "langroid[doc-chat,db]"
    ```
- To simply install _all_ optional dependencies, use the `all` extra (but note that this will result in longer load/startup times and a larger install size):
    ```bash
    pip install "langroid[all]"
    ```

??? note "Optional Installs for using SQL Chat with a PostgreSQL DB"
    If you are using `SQLChatAgent`
    (e.g. the script [`examples/data-qa/sql-chat/sql_chat.py`](https://github.com/langroid/langroid/blob/main/examples/data-qa/sql-chat/sql_chat.py),
    with a postgres db, you will need to:
    
    - Install PostgreSQL dev libraries for your platform, e.g.
        - `sudo apt-get install libpq-dev` on Ubuntu,
        - `brew install postgresql` on Mac, etc.
    - Install langroid with the postgres extra, e.g. `pip install langroid[postgres]`
      or `uv add "langroid[postgres]"` or `uv pip install --extra postgres -r pyproject.toml`.
      If this gives you an error, try 
      `uv pip install psycopg2-binary` in your virtualenv.


!!! tip "Work in a nice terminal, such as Iterm2, rather than a notebook"
    All of the examples we will go through are command-line applications.
    For the best experience we recommend you work in a nice terminal that supports 
    colored outputs, such as [Iterm2](https://iterm2.com/).    


!!! note "mysqlclient errors"
    If you get strange errors involving `mysqlclient`, try doing `pip uninstall mysqlclient` followed by `pip install mysqlclient` 

## Set up tokens/keys 

To get started, all you need is an OpenAI API Key.
If you don't have one, see [this OpenAI Page](https://platform.openai.com/docs/quickstart).
(Note that while this is the simplest way to get started, Langroid works with practically any LLM, not just those from OpenAI.
See the guides to using [Open/Local LLMs](https://langroid.github.io/langroid/tutorials/local-llm-setup/),
and other [non-OpenAI](https://langroid.github.io/langroid/tutorials/non-openai-llms/) proprietary LLMs.)

In the root of the repo, copy the `.env-template` file to a new file `.env`:
```bash
cp .env-template .env
```
Then insert your OpenAI API Key.
Your `.env` file should look like this:
```bash
OPENAI_API_KEY=your-key-here-without-quotes
```

Alternatively, you can set this as an environment variable in your shell
(you will need to do this every time you open a new shell):
```bash
export OPENAI_API_KEY=your-key-here-without-quotes
```

All of the following environment variable settings are optional, and some are only needed
to use specific features (as noted below).

- **Qdrant** Vector Store API Key, URL. This is only required if you want to use Qdrant cloud.
  Langroid uses LanceDB as the default vector store in its `DocChatAgent` class (for RAG).
  Alternatively [Chroma](https://docs.trychroma.com/) is also currently supported.
  We use the local-storage version of Chroma, so there is no need for an API key.
- **Redis** Password, host, port: This is optional, and only needed to cache LLM API responses
  using Redis Cloud. Redis [offers](https://redis.com/try-free/) a free 30MB Redis account
  which is more than sufficient to try out Langroid and even beyond.
  If you don't set up these, Langroid will use a pure-python
  Redis in-memory cache via the [Fakeredis](https://fakeredis.readthedocs.io/en/latest/) library.
- **GitHub** Personal Access Token (required for apps that need to analyze git
  repos; token-based API calls are less rate-limited). See this
  [GitHub page](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens).
- **Google Custom Search API Credentials:** Only needed to enable an Agent to use the `GoogleSearchTool`.
  To use Google Search as an LLM Tool/Plugin/function-call,
  you'll need to set up
  [a Google API key](https://developers.google.com/custom-search/v1/introduction#identify_your_application_to_google_with_api_key),
  then [setup a Google Custom Search Engine (CSE) and get the CSE ID](https://developers.google.com/custom-search/docs/tutorial/creatingcse).
  (Documentation for these can be challenging, we suggest asking GPT4 for a step-by-step guide.)
  After obtaining these credentials, store them as values of
  `GOOGLE_API_KEY` and `GOOGLE_CSE_ID` in your `.env` file.
  Full documentation on using this (and other such "stateless" tools) is coming soon, but
  in the meantime take a peek at the test
  [`tests/main/test_web_search_tools.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_web_search_tools.py) to see how to use it.


If you add all of these optional variables, your `.env` file should look like this:
```bash
OPENAI_API_KEY=your-key-here-without-quotes
GITHUB_ACCESS_TOKEN=your-personal-access-token-no-quotes
CACHE_TYPE=redis
REDIS_PASSWORD=your-redis-password-no-quotes
REDIS_HOST=your-redis-hostname-no-quotes
REDIS_PORT=your-redis-port-no-quotes
QDRANT_API_KEY=your-key
QDRANT_API_URL=https://your.url.here:6333 # note port number must be included
GOOGLE_API_KEY=your-key
GOOGLE_CSE_ID=your-cse-id
```

### Microsoft Azure OpenAI setup[Optional]

This section applies only if you are using Microsoft Azure OpenAI.

When using Azure OpenAI, additional environment variables are required in the
`.env` file.
This page [Microsoft Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/chatgpt-quickstart?tabs=command-line&pivots=programming-language-python#environment-variables)
provides more information, and you can set each environment variable as follows:

- `AZURE_OPENAI_API_KEY`, from the value of `API_KEY`
- `AZURE_OPENAI_API_BASE` from the value of `ENDPOINT`, typically looks like `https://your_resource.openai.azure.com`.
- For `AZURE_OPENAI_API_VERSION`, you can use the default value in `.env-template`, and latest version can be found [here](https://learn.microsoft.com/en-us/azure/ai-services/openai/whats-new#azure-openai-chat-completion-general-availability-ga)
- `AZURE_OPENAI_DEPLOYMENT_NAME` is an OPTIONAL deployment name which may be 
   defined by the user during the model setup.
- `AZURE_OPENAI_CHAT_MODEL` Azure OpenAI allows specific model names when you select the model for your deployment. You need to put precisely the exact model name that was selected. For example, GPT-3.5 (should be `gpt-35-turbo-16k` or `gpt-35-turbo`) or GPT-4 (should be `gpt-4-32k` or `gpt-4`).
- `AZURE_OPENAI_MODEL_NAME` (Deprecated, use `AZURE_OPENAI_CHAT_MODEL` instead).
  
!!! note "For Azure-based models use `AzureConfig` instead of `OpenAIGPTConfig`"
    In most of the docs you will see that LLMs are configured using `OpenAIGPTConfig`.
    However if you want to use Azure-deployed models, you should replace `OpenAIGPTConfig` with `AzureConfig`. 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)


## Next steps

Now you should be ready to use Langroid!
As a next step, you may want to see how you can use Langroid to [interact 
directly with the LLM](llm-interaction.md) (OpenAI GPT models only for now).
</file>

<file path="docs/quick-start/three-agent-chat-num-router.md">
# Three-Agent Collaboration, with message Routing

!!! tip "Script in `langroid-examples`"
    A full working example for the material in this section is
    in the `three-agent-chat-num-router.py` script in the `langroid-examples` repo:
    [`examples/quick-start/three-agent-chat-num-router.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/three-agent-chat-num-router.py).

Let's change the number game from the [three agent chat example](three-agent-chat-num.md) slightly.
In that example, when the `even_agent`'s LLM receives an odd number,
it responds with `DO-NOT-KNOW`, and similarly for the `odd_agent` when it
receives an even number. The `step()` method of the `repeater_task`
considers `DO-NOT-KNOW` to be an _invalid_ response and _continues_ to 
look for a valid response from any remaining sub-tasks.
Thus there was no need for the `processor_agent` to specify who should handle
the current number.

But what if there is a scenario where the `even_agent` and `odd_agent`
might return a legit but "wrong" answer?
In this section we add this twist -- when
the `even_agent` receives an odd number, it responds with -10, and similarly
for the `odd_agent` when it receives an even number.
We tell the `processor_agent` to avoid getting a negative number.

The goal we have set for the `processor_agent` implies that it 
must specify the intended recipient of 
the number it is sending. 
We can enforce this using a special Langroid Tool, 
[`RecipientTool`][langroid.agent.tools.recipient_tool.RecipientTool].
So when setting up the
`processor_task` we include instructions to use this tool
(whose name is `recipient_message`, the value of `RecipientTool.request`):

```py
processor_agent = lr.ChatAgent(config)
processor_task = lr.Task(
    processor_agent,
    name = "Processor",
    system_message="""
        You will receive a list of numbers from me (the user).
        Your goal is to apply a transformation to each number.
        However you do not know how to do this transformation.
        You can take the help of two people to perform the 
        transformation.
        If the number is even, send it to EvenHandler,
        and if it is odd, send it to OddHandler.
        
        IMPORTANT: send the numbers ONE AT A TIME
        
        The handlers will transform the number and give you a new number.        
        If you send it to the wrong person, you will receive a negative value.
        Your aim is to never get a negative number, so you must 
        clearly specify who you are sending the number to, using the
        `recipient_message` tool/function-call, where the `content` field
        is the number you want to send, and the `recipient` field is the name
        of the intended recipient, either "EvenHandler" or "OddHandler".        
        
        Once all numbers in the given list have been transformed, 
        say DONE and show me the result. 
        Start by asking me for the list of numbers.
    """,
    llm_delegate=True,
    single_round=False,
)
```

To enable the `processor_agent` to use this tool, we must enable it:
```py
processor_agent.enable_message(lr.agent.tools.RecipientTool)
```

The rest of the code remains the same as in the [previous section](three-agent-chat-num.md),
i.e., we simply add the two handler tasks
as sub-tasks of the `processor_task`, like this:
```python
processor_task.add_sub_task([even_task, odd_task])
```

One of the benefits of using the `RecipientTool` is that it contains 
mechanisms to remind the LLM to specify a recipient for its message,
when it forgets to do so (this does happen once in a while, even with GPT-4).


Feel free to try the working example script
`three-agent-chat-num-router.py` in the 
`langroid-examples` repo:
[`examples/quick-start/three-agent-chat-num-router.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/three-agent-chat-num-router.py):

```bash
python3 examples/quick-start/three-agent-chat-num-router.py
```

Below is screenshot of what this might look like, using the OpenAI function-calling 
mechanism with the `recipient_message` tool:

![three-agent-router-func.png](three-agent-router-func.png)

And here is what it looks like using Langroid's built-in tools mechanism (use the `-t` option when running the script):

![three-agent-router.png](three-agent-router.png)

And here is what it looks like using 
## Next steps

In the [next section](chat-agent-docs.md) you will learn
how to use Langroid with external documents.
</file>

<file path="docs/quick-start/three-agent-chat-num.md">
# Three-Agent Collaboration

!!! tip "Script in `langroid-examples`"
    A full working example for the material in this section is
    in the `three-agent-chat-num.py` script in the `langroid-examples` repo:
    [`examples/quick-start/three-agent-chat-num.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/three-agent-chat-num.py).


Let us set up a simple numbers exercise between 3 agents.
The `Processor` agent receives a number $n$, and its goal is to 
apply a transformation to the it. However it does not know how to apply the
transformation, and takes the help of two other agents to do so.
Given a number $n$,

- The `EvenHandler` returns $n/2$ if n is even, otherwise says `DO-NOT-KNOW`.
- The `OddHandler` returns $3n+1$ if n is odd, otherwise says `DO-NOT-KNOW`.

We'll first define a shared LLM config:

```py
llm_config = lr.language_models.OpenAIGPTConfig(
    chat_model=lr.language_models.OpenAIChatModel.GPT4o,
    # or, e.g., "ollama/qwen2.5-coder:latest", or "gemini/gemini-2.0-flash-exp"
)
```

Next define the config for the `Processor` agent:
```py
processor_config = lr.ChatAgentConfig(
    name="Processor",
    llm = llm_config,
    system_message="""
    You will receive a number from the user.
    Simply repeat that number, DO NOT SAY ANYTHING else,
    and wait for a TRANSFORMATION of the number 
    to be returned to you.
    
    Once you have received the RESULT, simply say "DONE",
    do not say anything else.
    """,        
    vecdb=None,
)
```

Then set up the `processor_agent`, along with the corresponding task:
```py
processor_agent = lr.ChatAgent(processor_config)

processor_task = lr.Task(
    processor_agent,
    llm_delegate=True, #(1)!
    interactive=False, #(2)!
    single_round=False, #(3)!
)

```

1. Setting the `llm_delegate` option to `True` means that the `processor_task` is
    delegated to the LLM (as opposed to the User), 
    in the sense that the LLM is the one "seeking" a response to the latest 
    number. Specifically, this means that in the `processor_task.step()` 
    when a sub-task returns `DO-NOT-KNOW`,
    it is _not_ considered a valid response, and the search for a valid response 
    continues to the next sub-task if any.
2. `interactive=False` means the task loop will not wait for user input.
3. `single_round=False` means that the `processor_task` should _not_ terminate after 
    a valid response from a responder.

Set up the other two agents and tasks:

```py
NO_ANSWER = lr.utils.constants.NO_ANSWER

even_config = lr.ChatAgentConfig(
    name="EvenHandler",
    llm = llm_config,
    system_message=f"""
    You will be given a number N. Respond as follows:
    
    - If N is even, divide N by 2 and show the result, 
      in the format: 
        RESULT = <result>
      and say NOTHING ELSE.
    - If N is odd, say {NO_ANSWER}
    """,    
)
even_agent = lr.ChatAgent(even_config)
even_task = lr.Task(
    even_agent,
    single_round=True,  # task done after 1 step() with valid response
)

odd_config = lr.ChatAgentConfig(
    name="OddHandler",
    llm = llm_config,
    system_message=f"""
    You will be given a number N. Respond as follows:
    
    - if N is odd, return the result (N*3+1), in the format:
        RESULT = <result> 
        and say NOTHING ELSE.
    
    - If N is even, say {NO_ANSWER}
    """,
)
odd_agent = lr.ChatAgent(odd_config)
odd_task = lr.Task(
    odd_agent,
    single_round=True,  # task done after 1 step() with valid response
)

```

Now add the `even_task` and `odd_task` as subtasks of the `processor_task`, 
and then run it with a number as input:

```python
processor_task.add_sub_task([even_task, odd_task])
processor_task.run(13)
```

The input number will be passed to the `Processor` agent as the user input.


Feel free to try the working example script
[`three-agent-chat-num.py`]()
`langroid-examples` repo:
[`examples/quick-start/three-agent-chat-num.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/three-agent-chat-num.py):

```bash
python3 examples/quick-start/three-agent-chat-num.py
```

Here's a screenshot of what it looks like:
![three-agent-num.png](three-agent-num.png)


## Next steps


In the [next section](chat-agent-tool.md) you will learn how to use Langroid
to equip a `ChatAgent` with tools or function-calling.
</file>

<file path="docs/quick-start/two-agent-chat-num.md">
# Two-Agent Collaboration

!!! tip "Script in `langroid-examples`"
    A full working example for the material in this section is
    in the `two-agent-chat-num.py` script in the `langroid-examples` repo:
    [`examples/quick-start/two-agent-chat-num.py`](https://github.com/langroid/langroid-examples/tree/main/examples/quick-start/two-agent-chat-num.py).


To illustrate these ideas, let's look at a toy example[^1] where 
a `Student` agent receives a list of numbers to add.
We set up this agent with an instruction that they do not know how to add,
and they can ask for help adding pairs of numbers.
To add pairs of numbers, we set up an `Adder` agent.

[^1]: Toy numerical examples are perfect to illustrate the ideas without
      incurring too much token cost from LLM API calls.

First define a common `llm_config` to use for both agents:
```python
llm_config = lr.language_models.OpenAIGPTConfig(
    chat_model=lr.language_models.OpenAIChatModel.GPT4o,
    # or, e.g., "ollama/qwen2.5-coder:latest", or "gemini/gemini-2.0-flash-exp"
)
```


Next, set up a config for the student agent, then create the agent
and the corresponding task:

```py
student_config = lr.ChatAgentConfig(
    name="Student",
    llm=llm_config,
    vecdb=None, #(1)!
    system_message="""
        You will receive a list of numbers from me (the User),
        and your goal is to calculate their sum.
        However you do not know how to add numbers.
        I can help you add numbers, two at a time, since
        I only know how to add pairs of numbers.
        Send me a pair of numbers to add, one at a time, 
        and I will tell you their sum.
        For each question, simply ask me the sum in math notation, 
        e.g., simply say "1 + 2", etc, and say nothing else.
        Once you have added all the numbers in the list, 
        say DONE and give me the final sum. 
        Start by asking me for the list of numbers.
    """,    
)
student_agent = lr.ChatAgent(student_config)
student_task = lr.Task(
    student_agent,
    name = "Student",
    llm_delegate = True, #(2)!
    single_round=False,  # (3)! 
)
```

1. We don't need access to external docs so we set `vecdb=None` to avoid 
   the overhead of loading a vector-store.
2. Whenever we "flip roles" and assign the LLM the role of generating questions, 
   we set `llm_delegate=True`. In effect this ensures that the LLM "decides" when
   the task is done.
3. This setting means the task is not a single-round task, i.e. it is _not_ done
   after one `step()` with a valid response.

Next, set up the Adder agent config, create the Adder agent
and the corresponding Task:

```py
adder_config = lr.ChatAgentConfig(
    name = "Adder", #(1)!
    llm=llm_config,
    vecdb=None,
    system_message="""
        You are an expert on addition of numbers. 
        When given numbers to add, simply return their sum, say nothing else
        """,     
)
adder_agent = lr.ChatAgent(adder_config)
adder_task = lr.Task(
    adder_agent,
    interactive=False, #(2)!
    single_round=True,  # task done after 1 step() with valid response (3)!
)
```
1. The Agent name is displayed in the conversation shown in the console.
2. Does not wait for user input.
3. We set `single_round=True` to ensure that the expert task is done after 
   one step() with a valid response. 

Finally, we add the `adder_task` as a sub-task of the `student_task`, 
and run the `student_task`:

```py
student_task.add_sub_task(adder_task) #(1)!
student_task.run()
```

1. When adding just one sub-task, we don't need to use a list.


For a full working example, see the 
[`two-agent-chat-num.py`](https://github.com/langroid/langroid-examples/blob/main/examples/quick-start/two-agent-chat-num.py)
script in the `langroid-examples` repo. You can run this using:
```bash
python3 examples/quick-start/two-agent-chat-num.py
```

Here is an example of the conversation that results:

![two-agent-num.png](two-agent-num.png)

## Logs of multi-agent interactions

!!! note "For advanced users"
    This section is for advanced users who want more visibility into the
    internals of multi-agent interactions.

When running a multi-agent chat, e.g. using `task.run()`, two types of logs
are generated:
- plain-text logs in `logs/<task_name>.log`
- tsv logs in `logs/<task_name>.tsv`

It is important to realize that the logs show _every iteration 
of the loop in `Task.step()`, i.e. every **attempt** at
responding to the current pending message, even those that are not allowed_.
The ones marked with an asterisk (*) are the ones that are considered valid
responses for a given `step()` (which is a "turn" in the conversation).

The plain text logs have color-coding ANSI chars to make them easier to read
by doing `less <log_file>`. The format is (subject to change):
```
(TaskName) Responder SenderEntity (EntityName) (=> Recipient) TOOL Content
```

The structure of the `tsv` logs is similar. A great way to view these is to
install and use the excellent `visidata` (https://www.visidata.org/) tool:
```bash
vd logs/<task_name>.tsv
```

## Next steps
As a next step, look at how to set up a collaboration among three agents
for a simple [numbers game](three-agent-chat-num.md).
</file>

<file path="docs/stylesheets/extra.css">
.md-logo img {
    height: 60px !important; /* Adjust size as necessary */
}
</file>

<file path="docs/tutorials/langroid-tour.md">
# A quick tour of Langroid

This is a quick tour of some Langroid features. For a more detailed guide,
see the [Getting Started guide](https://langroid.github.io/langroid/quick-start/).
There are many more features besides the ones shown here. To explore langroid more,
see the sections of the main [docs](https://langroid.github.io/langroid/),
and a 
[Colab notebook](https://colab.research.google.com/github/langroid/langroid/blob/main/examples/Langroid_quick_start.ipynb) 
you can try yourself.  


## Chat directly with LLM

Imports:

```python
import langroid as lr
import langroid.language_models as lm
```


Set up the LLM; note how you can specify the chat model -- if omitted, defaults
to OpenAI `GPT4o`. See the guide to using Langroid with 
[local/open LLMs](https://langroid.github.io/langroid/tutorials/local-llm-setup/),
and with [non-OpenAI LLMs](https://langroid.github.io/langroid/tutorials/non-openai-llms/).
    
```python
llm_config = lm.OpenAIGPTConfig( 
   chat_model="gpt-5-mini"
)
llm = lm.OpenAIGPT(llm_config)
```

Chat with bare LLM -- no chat accumulation, i.e. follow-up responses will *not*
be aware of prior conversation history (you need an Agent for that, see below).

```python
llm.chat("1 2 4 7 11 ?")
# ==> answers 16, with some explanation
```

## Agent

Make a [`ChatAgent`][langroid.agent.chat_agent.ChatAgent], 
and chat with it; now accumulates conv history

```python
agent = lr.ChatAgent(lr.ChatAgentConfig(llm=llm_config))
agent.llm_response("Find the next number: 1 2 4 7 11 ?")
# => responds 16
agent.llm_response("and then?)
# => answers 22
```

## Task

Make a [`Task`][langroid.agent.task.Task] and create a chat loop with the user:

```python
task = lr.Task(agent, interactive=True)
task.run()
```

## Tools/Functions/Structured outputs:

Define a [`ToolMessage`][langroid.agent.tool_message.ToolMessage] 
using Pydantic (v1) -- this gets transpiled into system-message instructions
to the LLM, so you never have to deal with writing a JSON schema.
(Besides JSON-based tools, Langroid also supports 
[XML-based tools](https://langroid.github.io/langroid/notes/xml-tools/), which 
are far more reliable when having the LLM return code in a structured output.)


```python
from pydantic import BaseModel

class CityTemperature(BaseModel):
    city: str
    temp: float

class WeatherTool(lr.ToolMessage):
    request: str = "weather_tool" #(1)!
    purpose: str = "To extract <city_temp> info from text" #(2)!

    city_temp: CityTemperature

    # tool handler
    def handle(self) -> CityTemperature:
        return self.city_temp
```

1. When this tool is enabled for an agent, a method named `weather_tool` gets auto-inserted in the agent class, 
   with body being the `handle` method -- this method handles the LLM's generation 
   of this tool.
2. The value of the `purpose` field is used to populate the system message to the LLM,
   along with the Tool's schema derived from its Pydantic-based definition.

Enable the Agent to use the `ToolMessage`, and set a system message describing the 
agent's task:

```python
agent.enable_message(WeatherTool)
agent.config.system_message = """
 Your job is to extract city and temperature info from user input
 and return it using the `weather_tool`.
"""
```

Create specialized task that returns a `CityTemperature` object:

```python
# configure task to terminate after (a) LLM emits a tool, (b) tool is handled by Agent
task_config = lr.TaskConfig(done_sequences=["T,A"])

# create a task that returns a CityTemperature object
task = lr.Task(agent, interactive=False, config=task_config)[CityTemperature]

# run task, with built-in tool-handling loop
data = task.run("It is 45 degrees F in Boston")

assert data.city == "Boston"
assert int(data.temp) == 45
```

## Chat with a document (RAG)

Create a [`DocChatAgent`][langroid.agent.special.doc_chat_agent.DocChatAgent].

```python
doc_agent_config = lr.agent.special.DocChatAgentConfig(llm=llm_config)
doc_agent = lr.agent.special.DocChatAgent(doc_agent_config)
```

Ingest the contents of a web page into the agent 
(this involves chunking, indexing into a vector-database, etc.):

```python
doc_agent.ingest_doc_paths("https://en.wikipedia.org/wiki/Ludwig_van_Beethoven")
```

Ask a question:

```
result = doc_agent.llm_response("When did Beethoven move from Bonn to Vienna?")
```

You should see the streamed response with citations like this:

![langroid-tour-beethoven.png](langroid-tour-beethoven.png)

## Two-agent interaction

Set up a teacher agent:

```python
from langroid.agent.tools.orchestration import DoneTool

teacher = lr.ChatAgent(
    lr.ChatAgentConfig(
        llm=llm_config,
        system_message=f"""
        Ask a numbers-based question, and your student will answer.
        You can then provide feedback or hints to the student to help them
        arrive at the right answer. Once you receive the right answer,
        use the `{DoneTool.name()}` tool to end the session.
        """
    )
)

teacher.enable_message(DoneTool)
teacher_task = lr.Task(teacher, interactive=False)

```

Set up a student agent:

```python
student = lr.ChatAgent(
    lr.ChatAgentConfig(
        llm=llm_config,
        system_message=f"""
        You will receive a numbers-related question. Answer to the best of
        your ability. If your answer is wrong, you will receive feedback or hints,
        and you can revise your answer, and repeat this process until you get 
        the right answer.
        """
    )
)

student_task = lr.Task(student, interactive=False, single_round=True)
```

Make the `student_task` a subtask of the `teacher_task`:

```python
teacher_task.add_sub_task(student_task)
```

Run the teacher task:

```python
teacher_task.run()
```

You should then see this type of interaction:

![langroid-tour-teacher.png](langroid-tour-teacher.png)
</file>

<file path="docs/tutorials/llm-usage-options.md">
# Options for accessing LLMs

> This is a work-in-progress document. It will be updated frequently.

The variety of ways to access the power of Large Language Models (LLMs) is growing 
rapidly, and there are a bewildering array of options. This document is an attempt to 
categorize and describe some of the most popular and useful ways to access LLMs,
via these 2x2x2  combinations:

- Websites (non-programmatic) or APIs (programmatic)
- Open-source or Proprietary 
- Chat-based interface or integrated assistive tools.

We will go into some of these combinations below. More will be added over time.

## Chat-based Web (non-API) access to Proprietary LLMs


This is best for *non-programmatic* use of LLMs: you go to a website and 
interact with the LLM via a chat interface -- 
you write prompts and/or upload documents, and the LLM responds with plain text
or can create artifacts (e.g. reports, code,
charts, podcasts, etc) that you can then copy into your files, workflow or codebase.
They typically allow you to upload text-based documents of various types, and some let you upload images, screen-shots, etc and ask questions about them.

Most of them are capable of doing *internet search* to inform their responses.


!!! note "Chat Interface vs Integrated Tools"
    Note that when using a chat-based interaction, you have to copy various artifacts
    from the web-site into another place, like your code editor, document, etc.
    AI-integrated tools relieve you of this burden by bringing the LLM power into 
    your workflow directly. More on this in a later section.

      
**Pre-requisites:** 

- *Computer*: Besides having a modern web browser (Chrome, Firefox, etc) and internet
access, there are no other special requirements, since the LLM is 
running on a remote server.
- *Coding knowledge*: Where (typically Python) code is produced, you will get best results
if you are conversant with Python so that you can understand and modify the code as
needed. In this category you do not need to know how to interact with an LLM API via code.

Here are some popular options in this category:

### OpenAI ChatGPT

Free access at [https://chatgpt.com/](https://chatgpt.com/)

With a ChatGPT-Plus monthly subscription ($20/month), you get additional features like:

- access to more powerful models
- access to [OpenAI canvas](https://help.openai.com/en/articles/9930697-what-is-the-canvas-feature-in-chatgpt-and-how-do-i-use-it) - this offers a richer interface than just a chat window, e.g. it automatically creates windows for code snippets, and shows results of running code
(e.g. output, charts etc).

Typical use: Since there is fixed monthly subscription (i.e. not metered by amount of 
usage), this is a cost-effective way to non-programmatically 
access a top LLM such as `GPT-4o` or `o1` 
(so-called "reasoning/thinking" models). Note however that there are limits on how many
queries you can make within a certain time period, but usually the limit is fairly
generous. 

What you can create, besides text-based artifacts:

- produce Python (or other language) code which you can copy/paste into notebooks or files
- SQL queries that you can copy/paste into a database tool
- Markdown-based tables
- You can't get diagrams, but you can get *code for diagrams*, 
e.g. python code for plots, [mermaid](https://github.com/mermaid-js/mermaid) code for flowcharts.
- images in some cases.

### OpenAI Custom GPTs (simply known as "GPTs")

[https://chatgpt.com/gpts/editor](https://chatgpt.com/gpts/editor)

Here you can conversationally interact with a "GPT Builder" that will 
create a version of ChatGPT
that is *customized* to your needs, i.e. with necessary background instructions,
context, and/or documents. 
The end result is a specialized GPT that you can then use for your specific
purpose and share with others (all of this is non-programmatic). 

E.g. [here](https://chatgpt.com/share/67153a4f-ea2c-8003-a6d3-cbc2412d78e5) is a "Knowledge Graph Builder" GPT

!!! note "Private GPTs requires an OpenAI Team Account"
    To share a custom GPT within a private group, you need an OpenAI Team account,
    see pricing [here](https://openai.com/chatgpt/pricing). Without a Team account,
    any shared GPT is public and can be accessed by anyone.


### Anthropic/Claude

[https://claude.ai](https://claude.ai)

The Claude basic web-based interface is similar to OpenAI ChatGPT, powered by 
Anthropic's proprietary LLMs. 
Anthropic's equivalent of ChatGPT-Plus is called "Claude Pro", which is also 
a $20/month subscription, giving you access to advanced models 
(e.g. `Claude-3.5-Sonnet`) and features.

Anthropic's equivalent of Custom GPTs is called 
[Projects](https://www.anthropic.com/news/projects), 
where you can create
an  LLM-powered interface that is augmented with your custom context and data.

Whichever product you are using, the interface auto-creates **artifacts** as needed --
these are stand-alone documents (code, text, images, web-pages, etc) 
that you may want to copy and paste into your own codebase, documents, etc.
For example you can prompt Claude to create full working interactive applications,
and copy the code, polish it and deploy it for others to use. See examples [here](https://simonwillison.net/2024/Oct/21/claude-artifacts/).

### Microsoft Copilot Lab

!!! note
    Microsoft's "Copilot" is an overloaded term that can refer to many different 
    AI-powered tools. Here we are referring to the one that is a collaboration between
    Microsoft and OpenAI, and is based on OpenAI's GPT-4o LLM, and powered by 
    Bing's search engine.

Accessible via [https://copilot.cloud.microsoft.com/](https://copilot.cloud.microsoft.com/)

The basic capabilities are similar to OpenAI's and Anthropic's offerings, but
come with so-called "enterprise grade" security and privacy features,
which purportedly make it suitable for use in educational and corporate settings.
Read more on what you can do with Copilot Lab [here](https://www.microsoft.com/en-us/microsoft-copilot/learn/?form=MA13FV).

Like the other proprietary offerings, Copilot can:

- perform internet search to inform its responses
- generate/run code and show results including charts

### Google Gemini

Accessible at [gemini.google.com](https://gemini.google.com).


## AI-powered productivity tools

These tools "bring the AI to your workflow", which is a massive productivity boost,
compared to repeatedly context-switching, e.g. copying/pasting between a chat-based AI web-app and your workflow.

- [**Cursor**](https://www.cursor.com/): AI Editor/Integrated Dev Environment (IDE). This is a fork of VSCode.
- [**Zed**](https://zed.dev/): built in Rust; can be customized to use Jetbrains/PyCharm keyboard shortcuts.
- [**Google Colab Notebooks with Gemini**](https://colab.research.google.com).
- [**Google NotebookLM**](https://notebooklm.google.com/): allows you to upload a set of text-based documents, 
  and create artifacts such as study guide, FAQ, summary, podcasts, etc.

    
## APIs for Proprietary LLMs

Using an API key allows *programmatic* access to the LLMs, meaning you can make
invocations to the LLM from within your own code, and receive back the results.
This is useful for building applications involving more complex workflows where LLMs
are used within a larger codebase, to access "intelligence" as needed.

E.g. suppose you are writing code that handles queries from a user, and you want to 
classify the user's _intent_ into one of 3 types: Information, or Action or Done.
Pre-LLMs, you would have had to write a bunch of rules or train a custom 
"intent classifier" that maps, for example:

- "What is the weather in Pittsburgh?" -> Information
- "Set a timer for 10 minutes" -> Action
- "Ok I have no more questions∞" -> Done

But using an LLM API, this is almost trivially easy - you instruct the LLM it should
classify the intent into one of these 3 types, and send the user query to the LLM,
and receive back the intent. 
(You can use Tools to make this robust, but that is outside the scope of this document.)

The most popular proprietary LLMs available via API are from OpenAI (or via  its
partner Microsoft), Anthropic, and Google:

- [OpenAI](https://platform.openai.com/docs/api-reference/introduction), to interact with `GPT-4o` family of models, and the `o1` family of "thinking/reasoning" models.
- [Anthropic](https://docs.anthropic.com/en/home) to use the `Claude` series of models.
- [Google](https://ai.google.dev/gemini-api/docs) to use the `Gemini` family of models.

These LLM providers are home to some of the most powerful LLMs available today,
specifically OpenAI's `GPT-4o` and Anthropic's `Claude-3.5-Sonnet`, and Google's `Gemini 1.5 Pro` (as of Oct 2024).

**Billing:** Unlike the fixed monthly subscriptions of ChatGPT, Claude and others, 
LLM usage via API is typically billed by *token usage*, i.e. you pay for the total
number of input and output "tokens" (a slightly technical term, but think of it as
a word for now).

Using an LLM API involves these steps:

- create an account on the provider's website as a "developer" or organization,
- get an API key,
- use the API key in your code to make requests to the LLM. 


**Prerequisites**:

- *Computer:* again, since the API is served over the internet, there are no special
  requirements for your computer.
- *Programming skills:* Using an LLM API involves either:
    - directly making REST API calls from your code, or 
    - use a scaffolding library (like [Langroid](https://github.com/langroid/langroid)) that abstracts away the details of the 
      API calls.
  
    In either case, you must be highly proficient in (Python) programming 
  to use this option.

## Web-interfaces to Open LLMs

!!! note  "Open LLMs"
    These are LLMs that have been publicly released, i.e. their parameters ("weights") 
    are publicly available -- we refer to these as *open-weight* LLMs. If in addition, the
    training datasets, and data-preprocessing and training code are also available, we would
    call these *open-source* LLMs. But lately there is a looser usage of the term "open-source",referring to just the weights being available. For our purposes we will just refer all of these models as **Open LLMs**.

There are many options here, but some popular ones are below. Note that some of these
are front-ends that allow you to interact with not only Open LLMs but also 
proprietary LLM APIs.

- [LMStudio](https://lmstudio.ai/)
- [OpenWebUI](https://github.com/open-webui/open-webui)
- [Msty](https://msty.app/)
- [AnythingLLM](https://anythingllm.com/)
- [LibreChat](https://www.librechat.ai/)


## API Access to Open LLMs

This is a good option if you are fairly proficient in (Python) coding. There are in 
fact two possibilities here:

- The LLM is hosted remotely, and you make REST API calls to the remote server. This
  is a good option when you want to run large LLMs and you don't have the resources (GPU and memory) to run them locally.
    - [groq](https://groq.com/) amazingly it is free, and you can run `llama-3.1-70b`
    - [cerebras](https://cerebras.ai/)
    - [open-router](https://openrouter.ai/)
- The LLM is running on your computer. This is a good option if your machine has sufficient RAM to accommodate the LLM you are trying to run, and if you are 
concerned about data privacy. The most user-friendly option is [Ollama](https://github.com/ollama/ollama); see more below.

Note that all of the above options provide an **OpenAI-Compatible API** to interact
with the LLM, which is a huge convenience: you can write code to interact with OpenAI's
LLMs (e.g. `GPT4o` etc) and then easily switch to one of the above options, typically
by changing a simple config (see the respective websites for instructions).

Of course, directly working with the raw LLM API quickly becomes tedious. This is where
a scaffolding library like [langroid](https://github.com/langroid/langroid) comes in
very handy - it abstracts away the details of the API calls, and provides a simple
programmatic interface to the LLM, and higher-level abstractions like 
Agents, Tasks, etc. Working with such a library is going to be far more productive
than directly working with the raw API. Below are instructions on how to use langroid
with some the above Open/Local LLM options.

See [here](https://langroid.github.io/langroid/tutorials/local-llm-setup/) for 
a guide to using Langroid with Open LLMs.
</file>

<file path="docs/auto_docstring.py">
from pathlib import Path
import mkdocs_gen_files

# -----------------------------------------------------#
#                    Configuration                    #
# -----------------------------------------------------#
src_dir = "langroid"
repo_root = "https://github.com/langroid/langroid/tree/main/"
nav = mkdocs_gen_files.Nav()

# -----------------------------------------------------#
#                       Runner                        #
# -----------------------------------------------------#
""" Generate code reference pages and navigation

    Based on the recipe of mkdocstrings:
    https://github.com/mkdocstrings/mkdocstrings

    Credits:
    Timothée Mazzucotelli
    https://github.com/pawamoy
"""
# Iterate over each Python file
for path in sorted(Path(src_dir).rglob("*.py")):
    if ".ipynb_checkpoints" in str(path):
        continue

    # Get path in module, documentation and absolute
    module_path = path.relative_to(src_dir).with_suffix("")
    doc_path = path.relative_to(src_dir).with_suffix(".md")
    full_doc_path = Path("reference", doc_path)

    # Handle edge cases
    parts = (src_dir,) + tuple(module_path.parts)
    if parts[-1] == "__init__":
        parts = parts[:-1]
        doc_path = doc_path.with_name("index.md")
        full_doc_path = full_doc_path.with_name("index.md")
    elif parts[-1] == "__main__":
        continue
    nav[parts] = doc_path.as_posix()

    # Write docstring documentation to disk via parser
    with mkdocs_gen_files.open(full_doc_path, "w") as fd:
        ident = ".".join(parts)
        full_code_path = repo_root + "/" + str(path)
        fd.write(f"[{path}]({full_code_path})\n")
        fd.write(f"::: {ident}")
    # Update parser
    mkdocs_gen_files.set_edit_path(full_doc_path, path)
    print(f"Doing docs for {full_doc_path}, {path}")

with mkdocs_gen_files.open("reference/SUMMARY.md", "w") as nav_file:
    nav_file.writelines(nav.build_literate_nav())
</file>

<file path="docs/FAQ.md">
# Frequently Asked Questions

## Can I view the reasoning (thinking) text when using a Reasoning LLM like R1 or o1?

Yes, see this note on [reasoning-content](https://langroid.github.io/langroid/notes/reasoning-content/).


## Does Langroid work with non-OpenAI LLMs?

Yes! Langroid works with practically any LLM, local or remote, closed or open.

See these two guides:

- [Using Langroid with local/open LLMs](https://langroid.github.io/langroid/tutorials/local-llm-setup/)
- [Using Langroid with non-OpenAI proprietary LLMs](https://langroid.github.io/langroid/tutorials/non-openai-llms/)

## Where can I find out about Langroid's architecture?

There are a few documents that can help:

- A work-in-progress [architecture description](https://langroid.github.io/langroid/blog/2024/08/15/overview-of-langroids-multi-agent-architecture-prelim/)
  on the Langroid blog.
- The Langroid [Getting Started](https://langroid.github.io/langroid/quick-start/) guide walks you 
  step-by-step through Langroid's features and architecture.
- An article by LanceDB on [Multi-Agent Programming with Langroid](https://lancedb.substack.com/p/langoid-multi-agent-programming-framework)

## How can I limit the number of output tokens generated by the LLM?

You can set the `max_output_tokens` parameter in the `LLMConfig` class,
or more commonly, the `OpenAIGPTConfig` class, which is a subclass of `LLMConfig`,
for example:

```python
import langroid as lr
import langroid.language_models as lm

llm_config = lm.OpenAIGPTConfig(
    chat_model="openai/gpt-3.5-turbo",
    max_output_tokens=100, # limit output to 100 tokens
)
agent_config = lr.ChatAgentConfig(
    llm=llm_config,
    # ... other configs
)
agent = lr.ChatAgent(agent_config)
```

Then every time the agent's `llm_response` method is called, the LLM's output 
will be limited to this number of tokens.

If you omit the `max_output_tokens`, it defaults to 8192. If you wish **not** to 
limit the output tokens, you can set `max_output_tokens=None`, in which case 
Langroid uses the model-specific maximum output tokens from the 
[`langroid/language_models/model_info.py`](https://github.com/langroid/langroid/blob/main/langroid/language_models/model_info.py) file
(specifically the `model_max_output_tokens` property of `LLMConfig`).
Note however that this model-specific may be quite large, so you would generally 
want to either omit setting `max_output_tokens` (which defaults to 8192), or set it
another desired value.


## How langroid handles long chat histories

You may encounter an error like this:

```
Error: Tried to shorten prompt history but ... longer than context length
```

This might happen when your chat history bumps against various limits.
Here is how Langroid handles long chat histories. Ultimately the LLM API is invoked with two key inputs:
the message history $h$, and the desired output length $n$ (defaults to the `max_output_tokens` in the 
`ChatAgentConfig`). These inputs are determined as follows (see the `ChatAgent._prep_llm_messages` method):

- let $H$ be the current message history, and $M$ be the value of `ChatAgentConfig.max_output_tokens`, and $C$ be 
  the context-length of the LLM.
- If $\text{tokens}(H) + M \leq C$, then langroid uses $h = H$ and $n = M$, since there is enough room to fit both the 
  actual chat history as well as the desired max output length.
- If $\text{tokens}(H) + M > C$, this means the context length is too small to accommodate the message history $H$ 
  and 
  the desired output length $M$. Then langroid tries to use a _shortened_ output length $n' = C - \text{tokens}(H)$, 
  i.e. the output is effectively _truncated_ to fit within the context length. 
    - If $n'$ is at least equal to `min_output_tokens` $m$ (default 10), langroid proceeds with $h = H$ and $n=n'$.
    - otherwise, this means that the message history $H$ is so long that the remaining space in the LLM's 
      context-length $C$ is unacceptably small (i.e. smaller than the minimum output length $m$). In this case,
      Langroid tries to shorten the message history by dropping early messages, and updating the message history $h$ as 
      long as $C - \text{tokens}(h) <  m$, until there are no more messages to drop (it will not drop the system 
      message or the last message, which is a user message), and throws the error mentioned above. 

If you are getting this error, you will want to check whether:

- you have set the `chat_context_length` too small, if you are setting it manually
- you have set the `max_output_tokens` too large
- you have set the `min_output_tokens` too large

If these look fine, then the next thing to look at is whether you are accumulating too much context into the agent 
history, for example retrieved passages (which can be very long) in a RAG scenario. One common case is when a query 
$Q$ is being answered using RAG, the retrieved passages $P$ are added to $Q$ to create a (potentially very long) prompt 
like 
> based on the passages P, answer query Q

Once the LLM returns an answer (if appropropriate for your context), you should avoid retaining the passages $P$ in the 
agent history, i.e. the last user message should be simply $Q$, rather than the prompt above. This functionality is exactly what you get when you 
use `ChatAgent._llm_response_temp_context`, which is used by default in the `DocChatAgent`. 

Another way to keep chat history tokens from growing too much is to use the `llm_response_forget` method, which 
erases both the query and response, if that makes sense in your scenario.

## How can I handle large results from Tools?

As of version 0.22.0, Langroid allows you to control the size of tool results
by setting [optional parameters](https://langroid.github.io/langroid/notes/large-tool-results/) 
in a `ToolMessage` definition.

## Can I handle a tool without running a task?

Yes, if you've enabled an agent to both _use_ (i.e. generate) and _handle_ a tool. 
See the `test_tool_no_task` for an example of this. The `NabroskiTool` is enabled
for the agent, and to get the agent's LLM to generate the tool, you first do 
something like:
```python
response = agent.llm_response("What is Nabroski of 1 and 2?")
```
Now the `response` is a `ChatDocument` that will contain the JSON for the `NabroskiTool`.
To _handle_ the tool, you will need to call the agent's `agent_response` method:

```python
result = agent.agent_response(response)
```

When you wrap the agent in a task object, and do `task.run()` the above two steps are done for you,
since Langroid operates via a loop mechanism, see docs 
[here](https://langroid.github.io/langroid/quick-start/multi-agent-task-delegation/#task-collaboration-via-sub-tasks).
The *advantage* of using `task.run()` instead of doing this yourself, is that this method
ensures that tool generation errors are sent back to the LLM so it retries the generation.

## OpenAI Tools and Function-calling support

Langroid supports OpenAI tool-calls API as well as OpenAI function-calls API.
Read more [here](https://github.com/langroid/langroid/releases/tag/0.7.0).

Langroid has always had its own native tool-calling support as well, 
which works with **any** LLM -- you can define a subclass of `ToolMessage` (pydantic based) 
and it is transpiled into system prompt instructions for the tool. 
In practice, we don't see much difference between using this vs OpenAI fn-calling. 
Example [here](https://github.com/langroid/langroid/blob/main/examples/basic/fn-call-local-simple.py).
Or search for `ToolMessage` in any of the `tests/` or `examples/` folders.

## Some example scripts appear to return to user input immediately without handling a tool.

This is because the `task` has been set up with `interactive=True` 
(which is the default). With this setting, the task loop waits for user input after
either the `llm_response` or `agent_response` (typically a tool-handling response) 
returns a valid response. If you want to progress through the task, you can simply 
hit return, unless the prompt indicates that the user needs to enter a response.

Alternatively, the `task` can be set up with `interactive=False` -- with this setting,
the task loop will _only_ wait for user input when an entity response (`llm_response` 
or `agent_response`) _explicitly_ addresses the user. Explicit user addressing can
be done using either:

- an orchestration tool, e.g. `SendTool` (see details in
the release notes for [0.9.0](https://github.com/langroid/langroid/releases/tag/0.9.0)), an example script is the [multi-agent-triage.py](https://github.com/langroid/langroid/blob/main/examples/basic/multi-agent-triage.py), or 
- a special addressing prefix, see the example script [1-agent-3-tools-address-user.py](https://github.com/langroid/langroid/blob/main/examples/basic/1-agent-3-tools-address-user.py)


## Can I specify top_k in OpenAIGPTConfig (for LLM API calls)?

No; Langroid currently only supports parameters accepted by OpenAI's API, and `top_k` is _not_ one of them. See:

- [OpenAI API Reference](https://platform.openai.com/docs/api-reference/chat/create)
- [Discussion on top_k, top_p, temperature](https://community.openai.com/t/temperature-top-p-and-top-k-for-chatbot-responses/295542/5)
- [Langroid example](https://github.com/langroid/langroid/blob/main/examples/basic/fn-call-local-numerical.py) showing how you can set other OpenAI API parameters, using the `OpenAICallParams` object.


## Can I persist agent state across multiple runs?

For example, you may want to stop the current python script, and 
run it again later, resuming your previous conversation.
Currently there is no built-in Langroid mechanism for this, but you can 
achieve a basic type of persistence by saving the agent's `message_history`:

-  if you used `Task.run()` in your script, make sure the task is 
set up with `restart=False` -- this prevents the agent state from being reset when 
the task is run again.
- using python's pickle module, you can save the `agent.message_history` to a file,
and load it (if it exists) at the start of your script.

See the example script [`chat-persist.py`](https://github.com/langroid/langroid/blob/main/examples/basic/chat-persist.py)

For more complex persistence, you can take advantage of the `GlobalState`,
where you can store message histories of multiple agents indexed by their name.
Simple examples of `GlobalState` are in the [`chat-tree.py`](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tree.py) example, 
and the [`test_global_state.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_global_state.py) test.

## Is it possible to share state between agents/tasks?

The above-mentioned `GlobalState` mechanism can be used to share state between 
agents/tasks. See the links mentioned in the previous answer.

## How can I suppress LLM output?

You can use the `quiet_mode` context manager for this, see 
[here](https://langroid.github.io/langroid/notes/quiet-mode/)

## How can I deal with LLMs (especially weak ones) generating bad JSON in tools?

Langroid already attempts to repair bad JSON (e.g. unescaped newlines, missing quotes, etc)  
using the [json-repair](https://github.com/mangiucugna/json_repair) library and other
custom methods, before attempting to parse it into a `ToolMessage` object.
However this type of repair may not be able to handle all edge cases of bad JSON 
from weak LLMs. There are two existing ways to deal with this, and one coming soon:

- If you are defining your own `ToolMessage` subclass, considering deriving it instead
  from `XMLToolMessage` instead, see the [XML-based Tools](https://langroid.github.io/langroid/notes/xml-tools/)
- If you are using an existing Langroid `ToolMessage`, e.g. `SendTool`, you can 
  define your own subclass of `SendTool`, say `XMLSendTool`,
  inheriting from both `SendTool` and `XMLToolMessage`; see this 
  [example](https://github.com/langroid/langroid/blob/main/examples/basic/xml_tool.py)
- Coming soon: strict decoding to leverage the Structured JSON outputs supported by OpenAI
  and open LLM providers such as `llama.cpp` and `vllm`.

The first two methods instruct the LLM to generate XML instead of JSON,
and any field that is designated with a `verbatim=True` will be enclosed 
within an XML `CDATA` tag, which does *not* require any escaping, and can
be far more reliable for tool-use than JSON, especially with weak LLMs.

## How can I handle an LLM "forgetting" to generate a `ToolMessage`? 

Sometimes the LLM (especially a weak one) forgets to generate a 
[`ToolMessage`][langroid.agent.tool_message.ToolMessage]
(either via OpenAI's tools/functions API, or via Langroid's JSON/XML Tool mechanism),
despite being instructed to do so. There are a few remedies Langroid offers for this:

**Improve the instructions in the `ToolMessage` definition:**

- Improve instructions in the `purpose` field of the `ToolMessage`.
- Add an `instructions` class-method to the `ToolMessage`, as in the
  [`chat-search.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/chat-search.py) script:

```python
@classmethod
def instructions(cls) -> str:
    return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """
```
  These instructions are meant to be general instructions on how to use the tool
  (e.g. how to set the field values), not to specifically about the formatting.

- Add a `format_instructions` class-method, e.g. like the one in the 
  [`chat-multi-extract-3.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/chat-multi-extract-3.py) 
  example script.

```python
@classmethod
def format_instructions(cls, tool: bool = True) -> str:
    instr = super().format_instructions(tool)
    instr += """
    ------------------------------
    ASK ME QUESTIONS ONE BY ONE, to FILL IN THE FIELDS 
    of the `lease_info` function/tool.
    First ask me for the start date of the lease.
    DO NOT ASK ANYTHING ELSE UNTIL YOU RECEIVE MY ANSWER.
    """
    return instr
```

**Override the `handle_message_fallback` method in the agent:**

This method is called when the Agent's `agent_response` method receives a non-tool
message as input. The default behavior of this method is to return None, but it
is very useful to override the method to handle cases where the LLM has forgotten
to use a tool. You can define this method to return a "nudge" to the LLM
telling it that it forgot to do a tool-call, e.g. see how it's done in the 
example script [`chat-multi-extract-local.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/chat-multi-extract-local.py):

```python
class LeasePresenterAgent(ChatAgent):
    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        """Handle scenario where Agent failed to present the Lease JSON"""
        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
            return """
            You either forgot to present the information in the JSON format
            required in `lease_info` JSON specification,
            or you may have used the wrong name of the tool or fields.
            Try again.
            """
        return None
```

Note that despite doing all of these, the LLM may still fail to generate a `ToolMessage`.
In such cases, you may want to consider using a better LLM, or an up-coming Langroid
feature that leverages **strict decoding** abilities of specific LLM providers
(e.g. OpenAI, llama.cpp, vllm) that are able to use grammar-constrained decoding
to force the output to conform to the specified structure.

Langroid also provides a simpler mechanism to specify the action to take
when an LLM does not generate a tool, via the `ChatAgentConfig.handle_llm_no_tool` 
config parameter, see the 
[docs](https://langroid.github.io/langroid/notes/handle-llm-no-tool/).

## Can I use Langroid to converse with a Knowledge Graph (KG)?

Yes, you can use Langroid to "chat with" either a Neo4j or ArangoDB KG, 
see docs [here](https://langroid.github.io/langroid/notes/knowledge-graphs/)

## How can I improve `DocChatAgent` (RAG) latency?

The behavior of `DocChatAgent` can be controlled by a number of settings in 
the `DocChatAgentConfig` class.
The top-level query-answering method in `DocChatAgent` is `llm_response`, which use the 
`answer_from_docs` method. At a high level, the response to an input message involves
the following steps:

- **Query to StandAlone:** LLM rephrases the query as a stand-alone query. 
   This can incur some latency. You can 
    turn it off by setting `assistant_mode=True` in the `DocChatAgentConfig`.
- **Retrieval:** The most relevant passages (chunks) are retrieved using a collection of semantic/lexical 
      similarity searches and ranking methods. There are various knobs in `DocChatAgentConfig` to control
      this retrieval.
- **Relevance Extraction:** LLM is used to retrieve verbatim relevant portions from
  the retrieved chunks. This is typically the biggest latency step. You can turn it off
  by setting the `relevance_extractor_config` to None in `DocChatAgentConfig`.
- **Answer Generation:** LLM generates answer based on retrieved passages.


See the [`doc-aware-chat.py`](https://github.com/langroid/langroid/blob/main/examples/docqa/doc-aware-chat.py)
example script, which illustrates some of these settings.

In some scenarios you want to *only* use the **retrieval** step of a `DocChatAgent`.
For this you can use the [`RetrievalTool`][langroid.agent.tools.retrieval_tool.RetrievalTool].
See the `test_retrieval_tool` in 
[`test_doc_chat_agent.py`](https://github.com/langroid/langroid/blob/main/tests/main/test_doc_chat_agent.py).
to learn how to use it. The above example script uses `RetrievalTool` as well.

## Is there support to run multiple tasks concurrently?

Yes, see the `run_batch_tasks` and related functions in 
[batch.py](https://github.com/langroid/langroid/blob/main/langroid/agent/batch.py).

See also:

- tests: [test_batch.py](https://github.com/langroid/langroid/blob/main/tests/main/test_batch.py),
   [test_relevance_extractor.py](https://github.com/langroid/langroid/blob/main/tests/main/test_relevance_extractor.py),
- example: [multi-agent-round-table.py](https://github.com/langroid/langroid/blob/main/examples/basic/multi-agent-round-table.py)

Another example is within 
[`DocChatAgent`](https://github.com/langroid/langroid/blob/main/langroid/agent/special/doc_chat_agent.py), 
which uses batch tasks for relevance extraction,
see the `get_verbatim_extracts` method -- when there are k relevant passages,
this runs k tasks concurrently, 
each of which uses an LLM-agent to extract relevant verbatim text from a passage.

## Can I use Langroid in a FastAPI server?

Yes, see the [langroid/fastapi-server](https://github.com/langroid/fastapi-server) repo.

## Can a sub-task end all parent tasks and return a result?

Yes, there are two ways to achieve this, using [`FinalResultTool`][langroid.agent.tools.orchestration.final_result_tool.FinalResultTool]:

From a `ChatAgent`'s tool-handler or `agent_response` method: Your code can return a 
`FinalResultTool` with arbitrary field types; this ends the current and all parent tasks and this  
`FinalResultTool` will appear as one of tools in the final `ChatDocument.tool_messages`.
See `test_tool_handlers_and_results` in 
[test_tool_messages.py](https://github.com/langroid/langroid/blob/main/tests/main/test_tool_messages.py), 
and [examples/basic/chat-tool-function.py](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tool-function.py)


From `ChatAgent`'s `llm_response` method: you can define a subclass of a 
`FinalResultTool` and enable the agent to use this tool, which means it will become
available for the LLM to generate. 
See [examples/basic/multi-agent-return-result.py](https://github.com/langroid/langroid/blob/main/examples/basic/multi-agent-return-result.py).

## How can I configure a task to retain or discard prior conversation?

In some scenarios, you may want to control whether each time you call a task's `run` 
method, the underlying agent retains the conversation history from the previous run.
There are two boolean config parameters that control this behavior: 

- the `restart` parameter (default `True`) in the `Task` constructor, and
- the `restart_as_subtask` (default `False`) parameter in the `TaskConfig` argument of the `Task` constructor.

To understand how these work, consider a simple scenario of a task `t` that has a 
subtask `t1`, e.g., suppose you have the following code with default settings 
of the `restart` and `restart_as_subtask` parameters:

```python
from langroid.agent.task import Task
from langroid.agent.task import TaskConfig

# default setttings:
rs = False
r = r1 = True

agent = ...
task_config = TaskConfig(restart_as_subtask=rs) 
t = Task(agent, restart=r, config=task_config)

agent1 = ...
t1 = Task(agent1, restart=r1, config=task_config)
t.add_subtask(t1)
```

This default setting works as follows:
Since task `t` was constructed with the default `restart=True`, when `t.run()` is called, the conversation histories of the agent underlying `t` as well as all 
those of all subtasks (such as `t1`) are reset. However, if during `t.run()`,
there are multiple calls to `t1.run()`, then the conversation history is retained across these calls, even though `t1` was constructed with the default `restart=True` --
this is because the `restart` constructor parameter has no effect on a task's reset
behavior **when it is a subtask**. 

The `TaskConfig.restart_as_subtask` parameter
controls the reset behavior of a task's `run` method when invoked as a subtask.
It defaults to `False`, which is why in the above example, the conversation history
of `t1` is retained across multiple calls to `t1.run()` that may occur
during execution of `t.run()`. If you set this parameter to `True` in the above
example, then the conversation history of `t1` would be reset each time `t1.run()` is called, during a call to `t.run()`.

To summarize, 

- The `Task` constructor's `restart` parameter controls the reset behavior of the task's `run` method when it is called directly, not as a subtask.
- The `TaskConfig.restart_as_subtask` parameter controls the reset behavior of the task's `run` method when it is called as a subtask.

These settings can be mixed and matched as needed.

Additionally, all reset behavior can be turned off during a specific `run()` invocation
by calling it with `allow_restart=False`, e.g.,  `t.run(..., allow_restart=False)`.

## How can I set up a task to exit as soon as the LLM responds?

In some cases you may want the top-level task or a subtask to exit as soon as the LLM responds. You can get this behavior by setting `single_round=True` during task construction, e.g.,

```python
from langroid.agent.task import Task

agent = ...
t = Task(agent, single_round=True, interactive=False)

result = t.run("What is 4 + 5?")
```

The name `single_round` comes from the fact that the task loop ends as soon as 
any **one** of the agent's responders return a valid response. Recall that an 
agent's responders are `llm_response`, `agent_response` (for tool handling), and `user_response` (for user input). In the above example there are no tools and no 
user interaction (since `interactive=False`), so the task will exit as soon as the LLM responds.

More commonly, you may only want this single-round behavior for a subtask, e.g.,

```python
agent = ...
t = Task(agent, single_round=False, interactive=True)

agent1 = ...
t1 = Task(agent1, single_round=True, interactive=False)

t.add_subtask(t1)
top_level_query = ...
result = t.run(...)
```

See the example script [`chat-2-agent-discuss.py`](https://github.com/langroid/langroid/blob/main/examples/basic/chat-2-agent-discuss.py) for an example of this, and also search for `single_round` in the rest of the examples.

!!! warning "Using `single_round=True` will prevent tool-handling"
    As explained above, setting `single_round=True` will cause the task to exit as soon as the LLM responds, and thus if it emits a valid tool (which the agent is enabled to handle), this tool will *not* be handled.
</file>

<file path="docs/index.md">
# Langroid: Harness LLMs with Multi-Agent Programming

## The LLM Opportunity

Given the remarkable abilities of recent Large Language Models (LLMs), there
is an unprecedented opportunity to build intelligent applications powered by
this transformative technology. The top question for any enterprise is: how
best to harness the power of LLMs for complex applications? For technical and
practical reasons, building LLM-powered applications is not as simple as
throwing a task at an LLM-system and expecting it to do it.

## Langroid's Multi-Agent Programming Framework

Effectively leveraging LLMs at scale requires a *principled programming 
framework*. In particular, there is often a need to maintain multiple LLM 
conversations, each instructed in different ways, and "responsible" for 
different aspects of a task.

An *agent* is a convenient abstraction that encapsulates LLM conversation 
state, along with access to long-term memory (vector-stores) and tools (a.k.a functions 
or plugins). Thus a **Multi-Agent Programming** framework is a natural fit 
for complex LLM-based applications.

> Langroid is the first Python LLM-application framework that was explicitly 
designed  with Agents as first-class citizens, and Multi-Agent Programming 
as the core  design principle. The framework is inspired by ideas from the 
[Actor Framework](https://en.wikipedia.org/wiki/Actor_model).

Langroid allows an intuitive definition of agents, tasks and task-delegation 
among agents. There is a principled mechanism to orchestrate multi-agent 
collaboration. Agents act as message-transformers, and take turns responding to (and
transforming) the current message. The architecture is lightweight, transparent, 
flexible, and allows other types of orchestration to be implemented; see the (WIP) 
[langroid architecture document](blog/posts/langroid-architecture.md).
Besides Agents, Langroid also provides simple ways to directly interact with LLMs and vector-stores. See the Langroid [quick-tour](tutorials/langroid-tour.md).

## Highlights
- **Agents as first-class citizens:** The `Agent` class encapsulates LLM conversation state,
  and optionally a vector-store and tools. Agents are a core abstraction in Langroid; 
  Agents act as _message transformers_, and by default provide 3 _responder_ methods, one corresponding to each 
  entity: LLM, Agent, User. 
- **Tasks:** A Task class wraps an Agent, gives the agent instructions (or roles, or goals),
  manages iteration over an Agent's responder methods,
  and orchestrates multi-agent interactions via hierarchical, recursive
  task-delegation. The `Task.run()` method has the same
  type-signature as an Agent's responder's methods, and this is key to how
  a task of an agent can delegate to other sub-tasks: from the point of view of a Task,
  sub-tasks are simply additional responders, to be used in a round-robin fashion
  after the agent's own responders.
- **Modularity, Reusabilily, Loose coupling:** The `Agent` and `Task` abstractions allow users to design
  Agents with specific skills, wrap them in Tasks, and combine tasks in a flexible way.
- **LLM Support**: Langroid works with practically any LLM, local/open or remote/proprietary/API-based, via a variety of libraries and providers. See guides to using [local LLMs](tutorials/local-llm-setup.md) and [non-OpenAI LLMs](tutorials/non-openai-llms.md). See [Supported LLMs](tutorials/supported-models.md).
- **Caching of LLM prompts, responses:** Langroid by default uses [Redis](https://redis.com/try-free/) for caching. 
- **Vector-stores**: [Qdrant](https://qdrant.tech/), [Chroma](https://www.trychroma.com/) and [LanceDB](https://www.lancedb.com/) are currently supported.
  Vector stores allow for Retrieval-Augmented-Generation (RAG).
- **Grounding and source-citation:** Access to external documents via vector-stores
  allows for grounding and source-citation.
- **Observability, Logging, Lineage:** Langroid generates detailed logs of multi-agent interactions and
  maintains provenance/lineage of messages, so that you can trace back
  the origin of a message.
- **Tools/Plugins/Function-calling**: Langroid supports OpenAI's recently
  released [function calling](https://platform.openai.com/docs/guides/gpt/function-calling)
  feature. In addition, Langroid has its own native equivalent, which we
  call **tools** (also known as "plugins" in other contexts). Function
  calling and tools have the same developer-facing interface, implemented
  using [Pydantic](https://docs.pydantic.dev/latest/),
  which makes it very easy to define tools/functions and enable agents
  to use them. Benefits of using Pydantic are that you never have to write
  complex JSON specs for function calling, and when the LLM
  hallucinates malformed JSON, the Pydantic error message is sent back to
  the LLM so it can fix it!



Don't worry if some of these terms are not clear to you. 
The [Getting Started Guide](quick-start/index.md) and subsequent pages 
will help you get up to speed.
</file>

<file path="examples/basic/multi-agent-search-critic/__init__.py">

</file>

<file path="examples/basic/multi-agent-search-critic/assistant_agent.py">
"""
AssistantAgent takes a user's question, breaks it down into smaller questions
for SearcherAgent to answer, and then presents the final answer; It then considers
feedback from CriticAgent, and may ask more questions or present the final answer
using a corrected reasoning.

Flow:

User Q ->
[L] -> QuestionTool(q1) ->
[A] -> validate, return QuestionTool(q1) ->
... AnswerTool(a1) from SearcherAgent ->
[A] -> AnswerTool(a1) -> natural lang ans for LLM
[L] -> either QuestionTool(q2) or FinalAnswerTool(steps, ans) ->
... if FinalAnswerTool(steps, ans) ->
[A] -> validate, return FinalAnswerTool(steps, ans) with recipient=Critic ->
... FeedbackTool(feedback, suggested_fix) from CriticAgent ->
[A] -> FeedbackTool(feedback, suggested_fix) -> natural lang feedback for LLM
[L] -> either QuestionTool(q2) or FinalAnswerTool(steps, ans) ->
...
"""

from typing import Optional

import typer

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool, PassTool

from .tools import AnswerTool, FeedbackTool, FinalAnswerTool, QuestionTool

app = typer.Typer()


class AssistantAgent(lr.ChatAgent):

    def init_state(self):
        super().init_state()
        self.expecting_question_tool: bool = False
        self.expecting_question_or_final_answer: bool = False  # expecting one of these
        # tools
        self.expecting_search_answer: bool = False
        self.original_query: str | None = None  # user's original query

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if self.expecting_question_or_final_answer:
            return f"""
            You may have intended to use a tool, but your JSON format may be wrong.
            
            REMINDER: You must do one of the following:
            - If you are ready with the final answer to the user's ORIGINAL QUERY
                [ Remember it was: {self.original_query} ],
              then present your reasoning steps and final answer using the 
              `final_answer_tool` in the specified JSON format.
            - If you still need to ask a question, then use the `question_tool`
              to ask a SINGLE question that can be answered from a web search.
            """
        elif self.expecting_question_tool:
            return f"""
            You must ask a question using the `question_tool` in the specified format,
            to break down the user's original query: {self.original_query} into 
            smaller questions that can be answered from a web search.
            """

    def question_tool(self, msg: QuestionTool) -> str | PassTool:
        self.expecting_search_answer = True
        self.expecting_question_tool = False
        # return the tool so it is handled by SearcherAgent
        # validated incoming, pass it on
        return PassTool()

    def answer_tool(self, msg: AnswerTool) -> str:
        self.expecting_question_or_final_answer = True
        self.expecting_search_answer = False
        return f"""
        Here is the answer to your question from the web search:
        {msg.answer}
        Now decide whether you want to:
        - present your FINAL answer to the user's ORIGINAL QUERY, OR
        - ask another question using the `question_tool`
            (Maybe REPHRASE the question to get BETTER search results).
        """

    def final_answer_tool(self, msg: FinalAnswerTool) -> ForwardTool | str:
        if not self.expecting_question_or_final_answer:
            return ""
        self.expecting_question_or_final_answer = False
        # insert the original query into the tool, in case LLM forgot to do so.
        msg.query = self.original_query
        # fwd to critic
        return ForwardTool(agent="Critic")

    def feedback_tool(self, msg: FeedbackTool) -> str:
        if msg.suggested_fix == "":
            return AgentDoneTool()
        else:
            self.expecting_question_or_final_answer = True
            # reset question count since feedback may initiate new questions
            return f"""
            Below is feedback about your answer. Take it into account to 
            improve your answer, EITHER by:
            - using the `final_answer_tool` again but with improved REASONING, OR
            - asking another question using the `question_tool`, and when you're 
                ready, present your final answer again using the `final_answer_tool`.
            
            FEEDBACK: {msg.feedback}
            SUGGESTED FIX: {msg.suggested_fix}
            """

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if self.original_query is None:
            self.original_query = (
                message if isinstance(message, str) else message.content
            )
            # just received user query, so we expect a question tool next
            self.expecting_question_tool = True

        if self.expecting_question_or_final_answer or self.expecting_question_tool:
            return super().llm_response(message)


def make_assistant_task(
    model: str,
    restart: bool = True,
) -> lr.Task:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )

    assistant_config = lr.ChatAgentConfig(
        system_message="""
        You are a resourceful assistant, able to think step by step to answer
        complex questions from the user. You must break down complex questions into
        simpler questions that can be answered by a web search. You must ask me 
        (the user) each question ONE BY ONE, using the `question_tool` in
         the specified format, and I will do a web search and send you
        a brief answer. Once you have enough information to answer my original
        (complex) question, you MUST present your INTERMEDIATE STEPS and FINAL ANSWER
        using the `final_answer_tool` in the specified JSON format.
        You will then receive FEEDBACK from the Critic, and if needed
        you should try to improve your answer based on this feedback,
        possibly by asking more questions.
        """,
        llm=llm_config,
        vecdb=None,
    )
    assistant_agent = AssistantAgent(assistant_config)
    assistant_agent.enable_message(QuestionTool)
    assistant_agent.enable_message(AnswerTool, use=False, handle=True)
    assistant_agent.enable_message(FinalAnswerTool)
    assistant_agent.enable_message(ForwardTool)
    assistant_agent.enable_message(PassTool)
    assistant_agent.enable_message(FeedbackTool, use=False, handle=True)

    assistant_task = lr.Task(
        assistant_agent,
        name="Assistant",
        llm_delegate=True,
        single_round=False,
        interactive=False,
        restart=restart,
    )

    return assistant_task
</file>

<file path="examples/basic/multi-agent-search-critic/critic_agent.py">
"""
CriticAgent task enforces FinalAnswerTool -> FeedbackTool, i.e.
- incoming msg must be a FinalAnswerTool
- outgoing msg must be a FeedbackTool

Flow:

FinalAnswerTool ->
[A] -> natural lang presentation to LLM
[L] -> FeedbackTool ->
[A] -> AgentDoneTool(FeedbackTool)

"""

import typer
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.configuration import Settings, set_global

from .tools import FeedbackTool, FinalAnswerTool

app = typer.Typer()


class CriticAgent(lr.ChatAgent):
    def init_state(self):
        super().init_state()
        self.expecting_feedback_tool: bool = False

    def final_answer_tool(self, msg: FinalAnswerTool) -> str:
        # received from Assistant. Extract the components as plain text,
        # so that the Critic LLM can provide feedback
        self.expecting_feedback_tool = True

        return f"""
        The user has presented the following query, intermediate steps and final answer
        shown below. Please provide feedback using the `feedback_tool`, 
        with the `feedback` field containing your feedback, and 
        the `suggested_fix` field containing a suggested fix, such as fixing how
        the answer or the steps, or how it was obtained from the steps, or 
        asking new questions.
        
        REMEMBER to set the `suggested_fix` field to an EMPTY string if the answer is 
        VALID.
        
        QUERY: {msg.query}
        
        STEPS: {msg.steps}
        
        ANSWER: {msg.answer}
        """

    def feedback_tool(self, msg: FeedbackTool) -> FeedbackTool:
        # validate, signal DONE, include the tool
        self.expecting_feedback_tool = False
        return AgentDoneTool(tools=[msg])

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if self.expecting_feedback_tool:
            return """
            You forgot to provide feedback using the `feedback_tool` 
            on the user's reasoning steps and final answer.
            """


def make_critic_task(model: str):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )
    critic_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message="""
        You excel at logical reasoning and combining pieces of information.
        You will receive a summary of the original query, intermediate steps and final 
        answer.
        You must examine these and provide feedback to the user, using the 
        `feedback_tool`, as follows:
        - If you think the answer and reasoning are valid, 
            simply set the `suggested_fix` field to an empty string "".
        - Otherwise set the `feedback` field to a reason why the answer is invalid,
            and in the `suggested_fix` field indicate how the user can improve the 
            answer, for example by reasoning differently, or asking different questions.
        """,
    )
    critic_agent = CriticAgent(critic_agent_config)
    critic_agent.enable_message(FeedbackTool)
    critic_agent.enable_message(FinalAnswerTool, use=False, handle=True)
    critic_task = lr.Task(
        critic_agent,
        name="Critic",
        interactive=False,
    )
    return critic_task


if __name__ == "__main__":

    @app.command()
    def main(
        debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
        model: str = typer.Option("", "--model", "-m", help="model name"),
        nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    ) -> None:
        set_global(
            Settings(
                debug=debug,
                cache=not nocache,
            )
        )
        load_dotenv()

        llm_config = lm.OpenAIGPTConfig(
            chat_model=model or lm.OpenAIChatModel.GPT4o,
            chat_context_length=16_000,
            temperature=0.2,
            max_output_tokens=500,
            timeout=45,
        )

        critic_agent_config = lr.ChatAgentConfig(
            llm=llm_config,
            vecdb=None,
            system_message="""
            You excel at logical reasoning and combining pieces of information.
            The user will send you a summary of the intermediate steps and final answer.
            You must examine these and provide feedback to the user, using the 
            `feedback_tool`, as follows:
            - If you think the answer and reasoning are valid, 
                simply set the `suggested_fix` field to an empty string "".
            - Otherwise set the `feedback` field to a reason why the answer is invalid,
                and in the `suggested_fix` field indicate how the user can improve the 
                answer, for example by reasoning differently, or asking different questions.
            """,
        )
        critic_agent = CriticAgent(critic_agent_config)
        critic_agent.enable_message(FeedbackTool)
        critic_agent.enable_message(FinalAnswerTool, use=False, handle=True)
        critic_task = lr.Task(
            critic_agent,
            name="Critic",
            interactive=False,
        )
        final_ans_tool = FinalAnswerTool(
            steps="""
            1. The moon landing was in 1969.
            2. Kennedy was president during 1969.            
            """,
            answer="Kennedy was president during the moon landing.",
        )
        # simulate receiving the tool from Assistant
        final_ans_doc = critic_agent.create_agent_response(
            tool_messages=[final_ans_tool]
        )
        result = critic_task.run(final_ans_doc)
        tools = critic_agent.get_tool_messages(result)
        assert len(tools) == 1
        assert isinstance(tools[0], FeedbackTool)

    app()
</file>

<file path="examples/basic/multi-agent-search-critic/main.py">
"""
Version of chat-search-assistant.py that is more likely to work local LLMs.

3-Agent system where:
- Assistant takes user's (complex) question, breaks it down into smaller pieces
    if needed
- Searcher takes Assistant's question, uses the Search tool to search the web
    (using DuckDuckGo), and returns a coherent answer to the Assistant.
- Critic takes Assistant's final answer, and provides feedback on it.

Once the Assistant thinks it has enough info to answer the user's question, it
says DONE and presents the answer to the user.

See also: chat-search for a basic single-agent search

Run like this from root of repo:

python3 -m examples.basic.multi-agent-search-critic.main

There are optional args, especially note these:

-m <model_name>: to run with a different LLM model (default: gpt4o)

For example try this question:

did Bach make more music than Beethoven?

You can specify a local LLM in a few different ways, e.g. `-m local/localhost:8000/v1`
or `-m ollama/mistral` etc. See here how to use Langroid with local LLMs:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

from langroid.utils.configuration import Settings, set_global

from .assistant_agent import make_assistant_task
from .critic_agent import make_critic_task
from .search_agent import make_search_task

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    print(
        """
        [blue]Welcome to the Web Search Assistant chatbot!
        I will try to answer your complex questions. 
        
        Enter x or q to quit at any point.
        """
    )
    load_dotenv()

    assistant_task = make_assistant_task(model)
    search_task = make_search_task(model)
    critic_task = make_critic_task(model)

    assistant_task.add_sub_task([search_task, critic_task])
    question = Prompt.ask("What do you want to know?")
    assistant_task.run(question)


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/multi-agent-search-critic/search_agent.py">
"""
SearcherAgent flow:

[A] stands for Agent response (i.e. agent_response)
[L] stands for LLM response (i.e. llm_response)

QuestionTool ->
[A] -> natural lang question for LLM ->
[L] -> DuckduckgoSearchTool ->
[A] -> results ->
[L] -> AnswerTool(results) ->
[A] -> AgentDoneTool(AnswerTool)

Note that this Agent's task enforces QuestionTool -> AnswerTool, i.e.
- incoming msg must be a QuestionTool
- outgoing msg must be an AnswerTool
"""

from typing import Optional

import typer
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.duckduckgo_search_tool import DuckduckgoSearchTool
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.configuration import Settings, set_global

from .tools import AnswerTool, QuestionTool

app = typer.Typer()

# class MyDDGSearchTool(DuckduckgoSearchTool):
#     request = "my_ddg_search"


class SearcherAgent(lr.ChatAgent):
    def init_state(self):
        super().init_state()
        self.curr_query: str | None = None
        self.expecting_search_results: bool = False
        self.expecting_search_tool: bool = False

    def __init__(self, config: lr.ChatAgentConfig):
        super().__init__(config)
        self.config = config
        self.enable_message(MetaphorSearchTool)  # DuckduckgoSearchTool
        self.enable_message(QuestionTool, use=False, handle=True)
        # agent is producing AnswerTool, so LLM should not be allowed to "use" it
        self.enable_message(AnswerTool, use=False, handle=True)

    def duckduckgo_search(self, msg: DuckduckgoSearchTool) -> str:
        """Override the DDG handler to update state"""
        self.expecting_search_results = True
        self.expecting_search_tool = False
        return msg.handle()

    def metaphor_search(self, msg: MetaphorSearchTool) -> str:
        """Override the Metaphor handler to update state"""
        self.expecting_search_results = True
        self.expecting_search_tool = False
        return msg.handle()

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        # we're here because msg has no tools
        if self.curr_query is None:
            # did not receive a question tool, so short-circuit and return None
            return None
        if self.expecting_search_tool:
            search_tool_name = MetaphorSearchTool.default_value("request")
            return f"""
            You forgot to use the web search tool`{search_tool_name}`  
            to answer the user's question : {self.curr_query}!!
            REMEMBER - you must ONLY answer the user's questions based on 
             results from a web-search, and you MUST NOT ANSWER them yourself.
             
            Please use the `{search_tool_name}` tool 
            using the specified JSON format, then compose your answer based on 
            the results from this web-search tool.
            """

    def question_tool(self, msg: QuestionTool) -> str:
        self.curr_query = msg.question
        self.expecting_search_tool = True
        search_tool_name = MetaphorSearchTool.default_value("request")
        return f"""
        User asked this question: {msg.question}.
        Perform a web search using the `{search_tool_name}` tool
        using the specified JSON format, to find the answer.
        """

    def answer_tool(self, msg: AnswerTool) -> AgentDoneTool:
        # signal DONE, and return the AnswerTool
        return AgentDoneTool(tools=[msg])

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if self.expecting_search_results:
            # message must be search results from the web search tool,
            # so let the LLM compose a response based on the search results

            curr_query = self.curr_query
            # reset state
            self.curr_query = None
            self.expecting_search_results = False
            self.expecting_search_tool = False

            result = super().llm_response_forget(message)

            # return an AnswerTool containing the answer,
            # with a nudge meant for the Assistant
            answer = f"""
                Here are the web-search results for the question: {curr_query}.
                ===
                {result.content}
                """

            ans_tool = AnswerTool(answer=answer)
            # cannot return a tool, so use this to create a ChatDocument
            return self.create_llm_response(tool_messages=[ans_tool])

        # Handling query from user (or other agent) => expecting a search tool
        result = super().llm_response_forget(message)
        return result


def make_search_task(model: str):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )

    search_tool_handler_method = MetaphorSearchTool.default_value("request")

    search_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message=f"""
        You are a web-searcher. For ANY question you get, you must use the
        `{search_tool_handler_method}` tool/function-call to get up to 5 results.
        Once you receive the results, you must compose a CONCISE answer 
        based on the search results and present the answer in this format:
        ANSWER: [... your CONCISE answer here ...]
        SOURCES: [links from the web-search that you used]
        
        EXTREMELY IMPORTANT: DO NOT MAKE UP ANSWERS, ONLY use the web-search results.
        """,
    )
    search_agent = SearcherAgent(search_agent_config)
    search_task = lr.Task(
        search_agent,
        name="Searcher",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    return search_task


if __name__ == "__main__":

    @app.command()
    def main(
        debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
        model: str = typer.Option("", "--model", "-m", help="model name"),
        nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    ) -> None:
        set_global(
            Settings(
                debug=debug,
                cache=not nocache,
            )
        )

        load_dotenv()

        search_task = make_search_task(model)
        # simulate an incoming message to this Task, from the Assistant agent
        q_doc = search_task.agent.create_agent_response(
            tool_messages=[QuestionTool(question="Who was Beethoven's teacher?")]
        )
        result = search_task.run(q_doc)
        tools = search_task.agent.get_tool_messages(result)
        assert len(tools) == 1
        assert isinstance(tools[0], AnswerTool)

    app()
</file>

<file path="examples/basic/multi-agent-search-critic/tools.py">
from typing import List

import typer

import langroid as lr

app = typer.Typer()


class QuestionTool(lr.ToolMessage):
    request: str = "question_tool"
    purpose: str = "Ask a SINGLE <question> that can be answered from a web search."
    question: str

    @classmethod
    def examples(cls) -> List[lr.ToolMessage]:
        return [
            cls(question="Which superconductor material was discovered in 2023?"),
            cls(question="What AI innovation did Meta achieve in 2024?"),
        ]


class AnswerTool(lr.ToolMessage):
    request: str = "answer_tool"
    purpose: str = "Present the <answer> to a web-search question"
    answer: str


class FinalAnswerTool(lr.ToolMessage):
    request: str = "final_answer_tool"
    purpose: str = """
        Present the intermediate <steps> and 
        final <answer> to the user's original <query>.
        """
    query: str
    steps: str
    answer: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            (
                "I want to show my reasoning steps, along with my final answer",
                cls(
                    query="was Plato mortal?",
                    steps="1. Man is mortal. 2. Plato was a man.",
                    answer="Plato was mortal.",
                ),
            ),
            cls(
                query="Who was president during the moon landing?",
                steps="1. The moon landing was in 1969. 2. Kennedy was president "
                "during 1969.",
                answer="Kennedy was president during the moon landing.",
            ),
        ]


class FeedbackTool(lr.ToolMessage):
    request: str = "feedback_tool"
    purpose: str = """
    Provide <feedback> on the user's answer. If the answer is valid based on the
    reasoning steps, then the feedback MUST be EMPTY
    """
    feedback: str
    suggested_fix: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            # just example
            cls(feedback="This looks fine!", suggested_fix=""),
            # thought + example
            (
                "I want to provide feedback on the reasoning steps and final answer",
                cls(
                    feedback="""
                    The answer is invalid because the conclusion does not follow from the
                    steps. Please check your reasoning and try again.
                    """,
                    suggested_fix="Check reasoning and try again",
                ),
            ),
        ]
</file>

<file path="examples/basic/multi-agent-search-critic-no-orch/__init__.py">

</file>

<file path="examples/basic/multi-agent-search-critic-no-orch/assistant_agent.py">
"""
AssistantAgent takes a user's question, breaks it down into smaller questions
for SearcherAgent to answer, and then presents the final answer; It then considers
feedback from CriticAgent, and may ask more questions or present the final answer
using a corrected reasoning.

Flow: (L stands for LLM, i.e. llm_response; A stands for Agent i.e. agent_response)

User Q ->
[L] -> QuestionTool(q1) ->
[A] -> validate, return QuestionTool(q1) ->
... AnswerTool(a1) from SearcherAgent ->
[A] -> AnswerTool(a1) -> natural lang ans for LLM
[L] -> either QuestionTool(q2) or FinalAnswerTool(steps, ans) ->
... if FinalAnswerTool(steps, ans) ->
[A] -> validate, return FinalAnswerTool(steps, ans) with recipient=Critic ->
... FeedbackTool(feedback, suggested_fix) from CriticAgent ->
[A] -> FeedbackTool(feedback, suggested_fix) -> natural lang feedback for LLM
[L] -> either QuestionTool(q2) or FinalAnswerTool(steps, ans) ->
...
"""

from typing import Optional

import typer

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool, PassTool

from .tools import AnswerTool, FeedbackTool, FinalAnswerTool, QuestionTool

app = typer.Typer()


class AssistantAgent(lr.ChatAgent):
    def init_state(self):
        super().init_state()
        self.expecting_question_tool: bool = False
        self.expecting_question_or_final_answer: bool = False  # expecting one of these
        # tools
        self.expecting_search_answer: bool = False
        self.original_query: str | None = None  # user's original query

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if self.expecting_question_or_final_answer:
            return f"""
            You may have intended to use a tool, but your JSON format may be wrong.
            
            REMINDER: You must do one of the following:
            - If you are ready with the final answer to the user's ORIGINAL QUERY
                [ Remember it was: {self.original_query} ],
              then present your reasoning steps and final answer using the 
              `final_answer_tool` in the specified JSON format.
            - If you still need to ask a question, then use the `question_tool`
              to ask a SINGLE question that can be answered from a web search.
            """
        elif self.expecting_question_tool:
            return f"""
            You must ask a question using the `question_tool` in the specified format,
            to break down the user's original query: {self.original_query} into 
            smaller questions that can be answered from a web search.
            """

    def question_tool(self, msg: QuestionTool) -> str | PassTool:
        self.expecting_search_answer = True
        self.expecting_question_tool = False
        # return the tool so it is handled by SearcherAgent
        # validated incoming, pass it on
        return AgentDoneTool(tools=[msg])

    def answer_tool(self, msg: AnswerTool) -> str:
        self.expecting_question_or_final_answer = True
        self.expecting_search_answer = False
        return f"""
        Here is the answer to your question from the web search:
        {msg.answer}
        Now decide whether you want to:
        - present your FINAL answer to the user's ORIGINAL QUERY, OR
        - ask another question using the `question_tool`
            (Maybe REPHRASE the question to get BETTER search results).
        """

    def final_answer_tool(self, msg: FinalAnswerTool) -> ForwardTool | str:
        if not self.expecting_question_or_final_answer:
            return ""
        self.expecting_question_or_final_answer = False
        # insert the original query into the tool, in case LLM forgot to do so.
        msg.query = self.original_query
        # fwd to critic
        return AgentDoneTool(tools=[msg])

    def feedback_tool(self, msg: FeedbackTool) -> str:
        if msg.suggested_fix == "":
            return AgentDoneTool()
        else:
            self.expecting_question_or_final_answer = True
            # reset question count since feedback may initiate new questions
            return f"""
            Below is feedback about your answer. Take it into account to 
            improve your answer, EITHER by:
            - using the `final_answer_tool` again but with improved REASONING, OR
            - asking another question using the `question_tool`, and when you're 
                ready, present your final answer again using the `final_answer_tool`.
            
            FEEDBACK: {msg.feedback}
            SUGGESTED FIX: {msg.suggested_fix}
            """

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if self.original_query is None:
            self.original_query = (
                message if isinstance(message, str) else message.content
            )
            # just received user query, so we expect a question tool next
            self.expecting_question_tool = True

        if self.expecting_question_or_final_answer or self.expecting_question_tool:
            return super().llm_response(message)


def make_assistant_task(
    model: str = "",
    restart: bool = True,
) -> lr.Task:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )

    assistant_config = lr.ChatAgentConfig(
        system_message="""
        You are a resourceful assistant, able to think step by step to answer
        complex questions from the user. You must break down complex questions into
        simpler questions that can be answered by a web search. You must ask me 
        (the user) each question ONE BY ONE, using the `question_tool` in
         the specified format, and I will do a web search and send you
        a brief answer. Once you have enough information to answer my original
        (complex) question, you MUST present your INTERMEDIATE STEPS and FINAL ANSWER
        using the `final_answer_tool` in the specified JSON format.
        You will then receive FEEDBACK from the Critic, and if needed
        you should try to improve your answer based on this feedback,
        possibly by asking more questions.
        """,
        llm=llm_config,
        vecdb=None,
    )
    assistant_agent = AssistantAgent(assistant_config)
    assistant_agent.enable_message(QuestionTool)
    assistant_agent.enable_message(AnswerTool, use=False, handle=True)
    assistant_agent.enable_message(FinalAnswerTool)
    assistant_agent.enable_message(ForwardTool)
    assistant_agent.enable_message(PassTool)
    assistant_agent.enable_message(FeedbackTool, use=False, handle=True)

    assistant_task = lr.Task(
        assistant_agent,
        name="Assistant",
        llm_delegate=True,
        single_round=False,
        interactive=False,
        restart=restart,
    )

    return assistant_task


if __name__ == "__main__":
    # restart = False, to preserve state across task.run() calls
    task = make_assistant_task(restart=False)
    question = task.run("which planet has more moons, Jupiter or Saturn?")
    assert isinstance(question.tool_messages[0], QuestionTool)
</file>

<file path="examples/basic/multi-agent-search-critic-no-orch/critic_agent.py">
"""
CriticAgent task enforces FinalAnswerTool -> FeedbackTool, i.e.
- incoming msg must be a FinalAnswerTool
- outgoing msg must be a FeedbackTool

Flow:

FinalAnswerTool ->
[A] -> natural lang presentation to LLM
[L] -> FeedbackTool ->
[A] -> AgentDoneTool(FeedbackTool)

"""

import typer
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.configuration import Settings, set_global

from .tools import FeedbackTool, FinalAnswerTool

app = typer.Typer()


class CriticAgent(lr.ChatAgent):
    def init_state(self):
        super().init_state()
        self.expecting_feedback_tool: bool = False

    def final_answer_tool(self, msg: FinalAnswerTool) -> str:
        # received from Assistant. Extract the components as plain text,
        # so that the Critic LLM can provide feedback
        self.expecting_feedback_tool = True

        return f"""
        The user has presented the following query, intermediate steps and final answer
        shown below. Please provide feedback using the `feedback_tool`, 
        with the `feedback` field containing your feedback, and 
        the `suggested_fix` field containing a suggested fix, such as fixing how
        the answer or the steps, or how it was obtained from the steps, or 
        asking new questions.
        
        REMEMBER to set the `suggested_fix` field to an EMPTY string if the answer is 
        VALID.
        
        QUERY: {msg.query}
        
        STEPS: {msg.steps}
        
        ANSWER: {msg.answer}
        """

    def feedback_tool(self, msg: FeedbackTool) -> FeedbackTool:
        # validate, signal DONE, include the tool
        self.expecting_feedback_tool = False
        return AgentDoneTool(tools=[msg])

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if self.expecting_feedback_tool:
            return """
            You forgot to provide feedback using the `feedback_tool` 
            on the user's reasoning steps and final answer.
            """


def make_critic_task(model: str):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )
    critic_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message="""
        You excel at logical reasoning and combining pieces of information.
        You will receive a summary of the original query, intermediate steps and final 
        answer.
        You must examine these and provide feedback to the user, using the 
        `feedback_tool`, as follows:
        - If you think the answer and reasoning are valid, 
            simply set the `suggested_fix` field to an empty string "".
        - Otherwise set the `feedback` field to a reason why the answer is invalid,
            and in the `suggested_fix` field indicate how the user can improve the 
            answer, for example by reasoning differently, or asking different questions.
        """,
    )
    critic_agent = CriticAgent(critic_agent_config)
    critic_agent.enable_message(FeedbackTool)
    critic_agent.enable_message(FinalAnswerTool, use=False, handle=True)
    critic_task = lr.Task(
        critic_agent,
        name="Critic",
        interactive=False,
    )
    return critic_task


if __name__ == "__main__":

    @app.command()
    def main(
        debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
        model: str = typer.Option("", "--model", "-m", help="model name"),
        nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    ) -> None:
        set_global(
            Settings(
                debug=debug,
                cache=not nocache,
            )
        )
        load_dotenv()

        llm_config = lm.OpenAIGPTConfig(
            chat_model=model or lm.OpenAIChatModel.GPT4o,
            chat_context_length=16_000,
            temperature=0.2,
            max_output_tokens=500,
            timeout=45,
        )

        critic_agent_config = lr.ChatAgentConfig(
            llm=llm_config,
            vecdb=None,
            system_message="""
            You excel at logical reasoning and combining pieces of information.
            The user will send you a summary of the intermediate steps and final answer.
            You must examine these and provide feedback to the user, using the 
            `feedback_tool`, as follows:
            - If you think the answer and reasoning are valid, 
                simply set the `suggested_fix` field to an empty string "".
            - Otherwise set the `feedback` field to a reason why the answer is invalid,
                and in the `suggested_fix` field indicate how the user can improve the 
                answer, for example by reasoning differently, or asking different questions.
            """,
        )
        critic_agent = CriticAgent(critic_agent_config)
        critic_agent.enable_message(FeedbackTool)
        critic_agent.enable_message(FinalAnswerTool, use=False, handle=True)
        critic_task = lr.Task(
            critic_agent,
            name="Critic",
            interactive=False,
        )
        final_ans_tool = FinalAnswerTool(
            steps="""
            1. The moon landing was in 1969.
            2. Kennedy was president during 1969.            
            """,
            answer="Kennedy was president during the moon landing.",
        )
        # simulate receiving the tool from Assistant
        final_ans_doc = critic_agent.create_agent_response(
            tool_messages=[final_ans_tool]
        )
        result = critic_task.run(final_ans_doc)
        tools = critic_agent.get_tool_messages(result)
        assert len(tools) == 1
        assert isinstance(tools[0], FeedbackTool)

    app()
</file>

<file path="examples/basic/multi-agent-search-critic-no-orch/main.py">
"""
Version of examples/basic/multi-agent-search-critic/main.py, but does NOT use any
inter-agent orchestration, i.e. we create a separate Task object from each agent,
but we do not connect them as sub-tasks.
Instead we write extra code to handle each task's output, and
determine what to do with it.

3-Agent system where:
- Assistant takes user's (complex) question, breaks it down into smaller pieces
    if needed
- Searcher takes Assistant's question, uses the Search tool to search the web
    (using DuckDuckGo), and returns a coherent answer to the Assistant.
- Critic takes Assistant's final answer, and provides feedback on it.

Once the Assistant thinks it has enough info to answer the user's question, it
says DONE and presents the answer to the user.

See also: chat-search for a basic single-agent search

Run like this from root of repo:

python3 -m examples.basic.multi-agent-search-critic-no-orch.main

There are optional args, especially note these:

-m <model_name>: to run with a different LLM model (default: gpt4o)

For example try this question:

did Bach make more music than Beethoven?

You can specify a local LLM in a few different ways, e.g. `-m local/localhost:8000/v1`
or `-m ollama/mistral` etc. See here how to use Langroid with local LLMs:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

import langroid as lr
from langroid.utils.configuration import Settings, set_global

from .assistant_agent import make_assistant_task
from .critic_agent import make_critic_task
from .search_agent import make_search_task
from .tools import AnswerTool, FeedbackTool, FinalAnswerTool, QuestionTool

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    print(
        """
        [blue]Welcome to the Web Search Assistant chatbot!
        I will try to answer your complex questions. 
        
        Enter x or q to quit at any point.
        """
    )
    load_dotenv()

    assistant_task = make_assistant_task(model, restart=False)
    search_task = make_search_task(model)
    critic_task = make_critic_task(model)

    def search_answer(qtool: QuestionTool) -> AnswerTool:
        """
        Take a QuestionTool, return an AnswerTool
        """
        return search_task[AnswerTool].run(qtool)

    def critic_feedback(fa: FinalAnswerTool) -> FeedbackTool:
        """
        Take a FinalAnswerTool, return a FeedbackTool
        """
        return critic_task[FeedbackTool].run(fa)

    def query_to_final_answer(question: str) -> FinalAnswerTool:
        """
        Take user's question, return FinalAnswerTool after
        iterating based on feedback from Critic.
        """

        question_tool_name = QuestionTool.default_value("request")
        final_answer_tool_name = FinalAnswerTool.default_value("request")

        tool = assistant_task[lr.ToolMessage].run(question)

        while True:
            if not isinstance(tool, (QuestionTool, FinalAnswerTool)):
                # no tool => nudge
                tool = assistant_task[lr.ToolMessage].run(
                    f"""
                     You forgot to use one of the tools:
                     `{question_tool_name}` or `{final_answer_tool_name}`.
                     """,
                )
            elif isinstance(tool, QuestionTool):
                # QuestionTool => get search result
                answer_tool = search_answer(tool)
                tool = assistant_task[lr.ToolMessage].run(answer_tool)
            else:
                # FinalAnswerTool => get feedback
                fb_tool = critic_feedback(tool)
                if fb_tool.suggested_fix == "":
                    # no suggested fix => return tool (which is a FinalAnswerTool)
                    return tool
                else:
                    # suggested fix => ask again
                    tool = assistant_task[lr.ToolMessage].run(fb_tool)

    # Interactive loop with user
    while True:
        question = Prompt.ask("What do you want to know?")
        if question.lower() in ["x", "q"]:
            break
        assistant_task.agent.init_state()
        final_answer = query_to_final_answer(question)
        assert isinstance(final_answer, FinalAnswerTool)


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/multi-agent-search-critic-no-orch/search_agent.py">
"""
SearcherAgent flow:

[A] stands for Agent response (i.e. agent_response)
[L] stands for LLM response (i.e. llm_response)

QuestionTool ->
[A] -> natural lang question for LLM ->
[L] -> DuckduckgoSearchTool ->
[A] -> results ->
[L] -> AnswerTool(results) ->
[A] -> AgentDoneTool(AnswerTool)

Note that this Agent's task enforces QuestionTool -> AnswerTool, i.e.
- incoming msg must be a QuestionTool
- outgoing msg must be an AnswerTool
"""

from typing import Optional

import typer
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.duckduckgo_search_tool import DuckduckgoSearchTool
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.configuration import Settings, set_global

from .tools import AnswerTool, QuestionTool

app = typer.Typer()

# class MyDDGSearchTool(DuckduckgoSearchTool):
#     request = "my_ddg_search"


class SearcherAgent(lr.ChatAgent):
    def init_state(self):
        super().init_state()
        self.curr_query: str | None = None
        self.expecting_search_results: bool = False
        self.expecting_search_tool: bool = False

    def __init__(self, config: lr.ChatAgentConfig):
        super().__init__(config)
        self.config = config
        self.enable_message(MetaphorSearchTool)  # DuckduckgoSearchTool
        self.enable_message(QuestionTool, use=False, handle=True)
        # agent is producing AnswerTool, so LLM should not be allowed to "use" it
        self.enable_message(AnswerTool, use=False, handle=True)

    def duckduckgo_search(self, msg: DuckduckgoSearchTool) -> str:
        """Override the DDG handler to update state"""
        self.expecting_search_results = True
        self.expecting_search_tool = False
        return msg.handle()

    def metaphor_search(self, msg: MetaphorSearchTool) -> str:
        """Override the Metaphor handler to update state"""
        self.expecting_search_results = True
        self.expecting_search_tool = False
        return msg.handle()

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        # we're here because msg has no tools
        if self.curr_query is None:
            # did not receive a question tool, so short-circuit and return None
            return None
        if self.expecting_search_tool:
            search_tool_name = MetaphorSearchTool.default_value("request")
            return f"""
            You forgot to use the web search tool`{search_tool_name}`  
            to answer the user's question : {self.curr_query}!!
            REMEMBER - you must ONLY answer the user's questions based on 
             results from a web-search, and you MUST NOT ANSWER them yourself.
             
            Please use the `{search_tool_name}` tool 
            using the specified JSON format, then compose your answer based on 
            the results from this web-search tool.
            """

    def question_tool(self, msg: QuestionTool) -> str:
        self.curr_query = msg.question
        self.expecting_search_tool = True
        search_tool_name = MetaphorSearchTool.default_value("request")
        return f"""
        User asked this question: {msg.question}.
        Perform a web search using the `{search_tool_name}` tool
        using the specified JSON format, to find the answer.
        """

    def answer_tool(self, msg: AnswerTool) -> AgentDoneTool:
        # signal DONE, and return the AnswerTool
        return AgentDoneTool(tools=[msg])

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if self.expecting_search_results:
            # message must be search results from the web search tool,
            # so let the LLM compose a response based on the search results

            curr_query = self.curr_query
            # reset state
            self.curr_query = None
            self.expecting_search_results = False
            self.expecting_search_tool = False

            result = super().llm_response_forget(message)

            # return an AnswerTool containing the answer,
            # with a nudge meant for the Assistant
            answer = f"""
                Here are the web-search results for the question: {curr_query}.
                ===
                {result.content}
                """

            ans_tool = AnswerTool(answer=answer)
            # cannot return a tool, so use this to create a ChatDocument
            return self.create_llm_response(tool_messages=[ans_tool])

        # Handling query from user (or other agent) => expecting a search tool
        result = super().llm_response_forget(message)
        return result


def make_search_task(model: str):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )

    search_tool_handler_method = MetaphorSearchTool.default_value("request")

    search_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message=f"""
        You are a web-searcher. For ANY question you get, you must use the
        `{search_tool_handler_method}` tool/function-call to get up to 5 results.
        Once you receive the results, you must compose a CONCISE answer 
        based on the search results and present the answer in this format:
        ANSWER: [... your CONCISE answer here ...]
        SOURCES: [links from the web-search that you used]
        
        EXTREMELY IMPORTANT: DO NOT MAKE UP ANSWERS, ONLY use the web-search results.
        """,
    )
    search_agent = SearcherAgent(search_agent_config)
    search_task = lr.Task(
        search_agent,
        name="Searcher",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    return search_task


if __name__ == "__main__":

    @app.command()
    def main(
        debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
        model: str = typer.Option("", "--model", "-m", help="model name"),
        nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    ) -> None:
        set_global(
            Settings(
                debug=debug,
                cache=not nocache,
            )
        )

        load_dotenv()

        search_task = make_search_task(model)
        # simulate an incoming message to this Task, from the Assistant agent
        q_doc = search_task.agent.create_agent_response(
            tool_messages=[QuestionTool(question="Who was Beethoven's teacher?")]
        )
        result = search_task.run(q_doc)
        tools = search_task.agent.get_tool_messages(result)
        assert len(tools) == 1
        assert isinstance(tools[0], AnswerTool)

    app()
</file>

<file path="examples/basic/multi-agent-search-critic-no-orch/tools.py">
from typing import List

import typer

import langroid as lr

app = typer.Typer()


class QuestionTool(lr.ToolMessage):
    request: str = "question_tool"
    purpose: str = "Ask a SINGLE <question> that can be answered from a web search."
    question: str

    @classmethod
    def examples(cls) -> List[lr.ToolMessage]:
        return [
            cls(question="Which superconductor material was discovered in 2023?"),
            cls(question="What AI innovation did Meta achieve in 2024?"),
        ]


class AnswerTool(lr.ToolMessage):
    request: str = "answer_tool"
    purpose: str = "Present the <answer> to a web-search question"
    answer: str


class FinalAnswerTool(lr.ToolMessage):
    request: str = "final_answer_tool"
    purpose: str = """
        Present the intermediate <steps> and 
        final <answer> to the user's original <query>.
        """
    query: str
    steps: str
    answer: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            (
                "I want to show my reasoning steps, along with my final answer",
                cls(
                    query="was Plato mortal?",
                    steps="1. Man is mortal. 2. Plato was a man.",
                    answer="Plato was mortal.",
                ),
            ),
            cls(
                query="Who was president during the moon landing?",
                steps="1. The moon landing was in 1969. 2. Kennedy was president "
                "during 1969.",
                answer="Kennedy was president during the moon landing.",
            ),
        ]


class FeedbackTool(lr.ToolMessage):
    request: str = "feedback_tool"
    purpose: str = """
    Provide <feedback> on the user's answer. If the answer is valid based on the
    reasoning steps, then the feedback MUST be EMPTY
    """
    feedback: str
    suggested_fix: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            # just example
            cls(feedback="This looks fine!", suggested_fix=""),
            # thought + example
            (
                "I want to provide feedback on the reasoning steps and final answer",
                cls(
                    feedback="""
                    The answer is invalid because the conclusion does not follow from the
                    steps. Please check your reasoning and try again.
                    """,
                    suggested_fix="Check reasoning and try again",
                ),
            ),
        ]
</file>

<file path="examples/basic/__init__.py">

</file>

<file path="examples/basic/1-agent-3-tools-address-user.py">
"""
Barebones example of a single agent using 3 tools.
Similar to 1-agent-3-tools.py, but here the task is set up
with `interactive=False`, meaning user input is awaited only
when user is explicitly addressed using an addressing prefix.
"""

from typing import Any, List, Tuple

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ForwardTool
from langroid.utils.configuration import settings
from langroid.utils.constants import AT

DEFAULT_LLM = lm.OpenAIChatModel.GPT4o

# (1) DEFINE THE TOOLS


class UpdateTool(lr.ToolMessage):
    request: str = "update"
    purpose: str = "To update the stored number to the given <number>"
    number: int

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        # Examples that will be compiled into few-shot examples for the LLM.
        # Each example can either be...
        return [
            cls(number=3),  # ... just instances of the tool-class, OR
            (  # ...a tuple of "thought leading to tool", and the tool instance
                "I want to update the stored number to number 4 from the user",
                cls(number=4),
            ),
        ]


class AddTool(lr.ToolMessage):
    request: str = "add"
    purpose: str = "To add the given <number> to the stored number"
    number: int

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        return [
            cls(number=3),
            (
                "I want to add number 10 to the stored number",
                cls(number=10),
            ),
        ]


class ShowTool(lr.ToolMessage):
    request: str = "show"
    purpose: str = "To show the user the stored <number>"

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        return [
            cls(number=3),
            (
                "I want to show the user the stored number 10",
                cls(number=10),
            ),
        ]


# (2) DEFINE THE AGENT, with the tool-handling methods
class NumberAgent(lr.ChatAgent):
    secret: int = 0

    def update(self, msg: UpdateTool) -> str:
        self.secret = msg.number
        return f"Ok I updated the stored number to {msg.number}"

    def add(self, msg: AddTool) -> str:
        self.secret += msg.number
        return f"Added {msg.number} to stored number => {self.secret}"

    def show(self, msg: ShowTool) -> str:
        return f"Inform the user that the SECRET NUMBER is {self.secret}"

    def handle_message_fallback(self, msg: str | lr.ChatDocument) -> Any:
        """
        If we're here it means there was no recognized tool in `msg`.
        So if it was from LLM, use ForwardTool to send to user.
        """
        if isinstance(msg, lr.ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            return ForwardTool(agent="User")


def app(
    m: str = DEFAULT_LLM,  # pass -d <model> to use non-default LLM
    d: bool = False,  # pass -d to enable debug mode (see prompts etc)
    nc: bool = False,  # pass -nc to disable cache-retrieval (i.e. get fresh answers)
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=4096,  # set this based on model
        max_output_tokens=100,
        temperature=0.2,
        stream=True,
        timeout=45,
    )

    # (3) CREATE THE AGENT
    agent_config = lr.ChatAgentConfig(
        name="NumberAgent",
        llm=llm_cfg,
        system_message=f"""
        When the user's request matches one of your available tools, use it, 
        otherwise respond directly to the user.
        NOTE: Whenever you want to address the user directly, you MUST
        use "{AT}User", followed by your message. 
        """,
    )

    agent = NumberAgent(agent_config)

    # (4) ENABLE/ATTACH THE TOOLS to the AGENT

    agent.enable_message(UpdateTool)
    agent.enable_message(AddTool)
    agent.enable_message(ShowTool)

    # (5) CREATE AND RUN THE TASK
    task_config = lr.TaskConfig(addressing_prefix=AT)
    task = lr.Task(agent, interactive=False, config=task_config)

    """
    Note: try saying these when it waits for user input:
    
    add 10
    update 50
    add 3
    show 
    """

    task.run()


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/basic/1-agent-3-tools.py">
"""
Barebones example of a single agent using 3 tools.

"""

from typing import Any, List, Tuple

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ForwardTool
from langroid.utils.configuration import settings

DEFAULT_LLM = lm.OpenAIChatModel.GPT4o

# (1) DEFINE THE TOOLS


class UpdateTool(lr.ToolMessage):
    request: str = "update"
    purpose: str = "To update the stored number to the given <number>"
    number: int

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        # Examples that will be compiled into few-shot examples for the LLM.
        # Each example can either be...
        return [
            cls(number=3),  # ... just instances of the tool-class, OR
            (  # ...a tuple of "thought leading to tool", and the tool instance
                "I want to update the stored number to number 4 from the user",
                cls(number=4),
            ),
        ]


class AddTool(lr.ToolMessage):
    request: str = "add"
    purpose: str = "To add the given <number> to the stored number"
    number: int

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        return [
            cls(number=3),
            (
                "I want to add number 10 to the stored number",
                cls(number=10),
            ),
        ]


class ShowTool(lr.ToolMessage):
    request: str = "show"
    purpose: str = "To show the user the stored <number>"

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        return [
            cls(number=3),
            (
                "I want to show the user the stored number 10",
                cls(number=10),
            ),
        ]


# (2) DEFINE THE AGENT, with the tool-handling methods
class NumberAgent(lr.ChatAgent):
    secret: int = 0

    def update(self, msg: UpdateTool) -> str:
        self.secret = msg.number
        return f"""
            Ok I updated the stored number to {msg.number}.
            Ask the user what they want to do
        """

    def add(self, msg: AddTool) -> str:
        self.secret += msg.number
        return f"""
            Added {msg.number} to stored number => {self.secret}.
            Ask the user what they want to do.
        """

    def show(self, msg: ShowTool) -> str:
        return f"Tell the user that the SECRET NUMBER is {self.secret}"

    def handle_message_fallback(self, msg: str | lr.ChatDocument) -> Any:
        """
        If we're here it means there was no recognized tool in `msg`.
        So if it was from LLM, use ForwardTool to send to user.
        """
        if isinstance(msg, lr.ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            return ForwardTool(agent="User")


def app(
    m: str = DEFAULT_LLM,  # model
    d: bool = False,  # pass -d to enable debug mode (see prompts etc)
    nc: bool = False,  # pass -nc to disable cache-retrieval (i.e. get fresh answers)
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=4096,  # set this based on model
        max_output_tokens=100,
        temperature=0.2,
        stream=True,
        timeout=45,
    )

    # (3) CREATE THE AGENT
    agent_config = lr.ChatAgentConfig(
        name="NumberAgent",
        llm=llm_cfg,
        system_message="""
        When the user's request matches one of your available tools, use it, 
        otherwise respond directly to the user.
        """,
    )

    agent = NumberAgent(agent_config)

    # (4) ENABLE/ATTACH THE TOOLS to the AGENT

    agent.enable_message(UpdateTool)
    agent.enable_message(AddTool)
    agent.enable_message(ShowTool)

    # (5) CREATE AND RUN THE TASK
    task = lr.Task(agent, interactive=False)

    """
    Note: try saying these when it waits for user input:
    
    add 10
    update 50
    add 3
    show <--- in this case remember to hit enter when it waits for your input.
    """
    task.run()


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/basic/1d-screen-click.py">
"""

A Bit-Shooter Game played on a 1-dimensional binary screen.

Given an LLM Agent access to a 1-dimensional "screen" represented
as a string of bits (0s and 1s), e.g. "101010",
and equip it with a "Click tool" (like a mouse click) that allows it to
click on a bit -- clicking the bit causes it to flip.

The Agent plays a "Bit Shooter" game where the goal is to get rid of all
1s in the "screen".

To use the Click tool, the Agent must specify the position (zero-based)
where it wants to click. This causes the bit to flip.
The LLM is then presented with the new state of the screen,
and the process repeats until all 1s are gone.

Clearly the Agent (LLM) needs to be able to accurately count the bit positions,
to be able to correctly click on the 1s.

Run like this (--model is optional, defaults to GPT4o):

python3 examples/basic/1d-screen-click.py --model litellm/anthropic/claude-3-5-sonnet-20241022

At the beginning you get to specify the initial state of the screen:
- size of the screen (how many bits)
- the (0-based) locations of the 1s (SPACE-separated) in the screen.

E.g. try this:
- size = 50,
- 1-indices: 0 20 30 40

The loop is set to run in interactive mode (to prevent runaway loops),
so you have to keep hitting enter to see the LLM's next move.

The main observation is that when you run it with claude-3.5-sonnet,
the accuracy of the Agent's clicks is far superior to other LLMs like GPT-4o
and even GPT-4.

To try with other LLMs, you can set the --model param to, for example:
- gpt-4 (set OPENAI_API_KEY in your env or .env file)
- gpt-4o (ditto, set OPENAI_API_KEY)
- groq/llama-3.1-70b-versatile (set GROQ_API_KEY in your env or .env file)
- cerebras/llama3.1-70b (set CEREBRAS_API_KEY in your env or .env file)
- ollama/qwen2.5-coder:latest

See here for a full guide on local/open LLM setup with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
And here for how to use with other non-OpenAPI LLMs:
https://langroid.github.io/langroid/tutorials/non-openai-llms/
"""

from typing import List, Tuple

import fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import AgentDoneTool
from pydantic import BaseModel
from langroid.utils.globals import GlobalState


class ScreenState(BaseModel):
    """
    Represents the state of the 1-dimensional binary screen
    """

    screen: str | None = None  # binary string, e.g. "101010"

    def __init__(
        self,
        one_indices: List[int] = [1],
        size: int = 1,
    ):
        super().__init__()
        # Initialize with all zeros
        screen_list = ["0"] * size

        # Set 1s at specified indices
        for idx in one_indices:
            if 0 <= idx < size:
                screen_list[idx] = "1"

        # Join into string
        self.screen = "".join(screen_list)

    @classmethod
    def set_state(
        cls,
        one_indices: List[int],
        size: int,
    ) -> "ScreenState":
        """
        Factory method to create and set initial state.
        """
        initial_state = cls(
            one_indices=one_indices,
            size=size,
        )
        GlobalScreenState.set_values(state=initial_state)

    def flip(self, i: int):
        """
        Flip the i-th bit
        """
        if self.screen is None or i < 0 or i >= len(self.screen):
            return

        screen_list = list(self.screen)
        screen_list[i] = "1" if screen_list[i] == "0" else "0"
        self.screen = "".join(screen_list)


class GlobalScreenState(GlobalState):
    state: ScreenState = ScreenState()


def get_state() -> ScreenState:
    return GlobalScreenState.get_value("state")


class ClickTool(lr.ToolMessage):
    request: str = "click_tool"
    purpose: str = """
        To click at <position> on the 1-dimensional binary screen, 
        which causes the bit at that position to FLIP.
        IMPORTANT: the position numbering starts from 0!!!
    """

    position: int

    @classmethod
    def examples(cls) -> List[lr.ToolMessage | Tuple[str, lr.ToolMessage]]:
        return [
            cls(position=3),
            (
                "I want to click at position 5",
                cls(position=5),
            ),
        ]

    def handle(self) -> str | AgentDoneTool:
        state = get_state()
        state.flip(self.position)
        print("SCREEN STATE = ", state.screen)
        if "1" not in state.screen:
            return AgentDoneTool()
        return state.screen


def main(model: str = ""):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    click_tool_name = ClickTool.default_value("request")
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            name="Clicker",
            llm=llm_config,
            use_functions_api=False,  # suppress OpenAI functions/tools
            use_tools=True,  # enable langroid-native tools: works with any LLM
            show_stats=False,
            system_message=f"""
            You are an expert at COMPUTER USE.
            In this task you only have to be able to understand a 1-dimensional 
            screen presented to you as a string of bits (0s and 1s).
            You will play a 1-dimensional BIT-shooter game!
            
            Your task is to CLICK ON THE LEFTMOST 1 in the bit-string, 
            to flip it to a 0.
            
            Always try to click on the LEFTMOST 1 in the bit-sequence. 
            
            To CLICK on the screen you 
            must use the TOOL `{click_tool_name}` where the  
            `position` field specifies the position (zero-based) to click.
            If you CORRECTLY click on a 1, the bit at that position will be 
            turned to 0.
            But if you click on a 0, it will turn into a 1, 
            taking you further from your goal.
            
            So you MUST ACCURATELY specify the position of the LEFTMOST 1 to click,
            making SURE there is a 1 at that position.
            In other words, it is critical that you are able to ACCURATELY COUNT 
            the bit positions so that you are able to correctly identify the position 
            of the LEFTMOST 1 bit in the "screen" given to you as a string of bits.
            """,
        )
    )

    agent.enable_message(ClickTool)

    task = lr.Task(agent, interactive=True, only_user_quits_root=False)

    # kick it off with initial screen state (set below by user)
    task.run(get_state())


if __name__ == "__main__":
    size = int(Prompt.ask("Size of screen (how many bits)"))
    ones = Prompt.ask("Indices of 1s (SPACE-separated)").split(" ")
    ones = [int(x) for x in ones]
    ScreenState.set_state(ones, size)
    print("SCREEN STATE = ", get_state().screen)
    fire.Fire(main)
</file>

<file path="examples/basic/2-agent-tools.py">
"""
2 Agent setup where Main agent asks a question, Helper has a few tools to help answer,
and for any question, Helper finishes after first use of any tool.

Run like this:

python3 examples/basic/2-agent-tools.py

When it waits for user input, try asking things like:

- capital of uganda?
    => Main answers
- polinsky of 4?
    => Main says do not know, handled by helper, who returns answer
- chichikov of 5?
    => Main says do not know, handled by helper, who returns answer
"""

from typing import Any

import langroid as lr
from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool


class MainChatAgent(lr.ChatAgent):
    def handle_message_fallback(self, msg: str | lr.ChatDocument) -> Any:
        """
        We'd be here if there were no recognized tools in the incoming msg.
        If this was from LLM, forward to user.
        """
        if isinstance(msg, lr.ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            return ForwardTool(agent="User")


main = MainChatAgent(
    lr.ChatAgentConfig(
        name="Main",
        system_message=f"""
        Help the user with their questions. When you don't know the answer, 
        simply say {lr.utils.constants.NO_ANSWER} and nothing else.
        Your Helper will attempt to handle the question, and send you back their
        answer, and you can present it to the user.   
        
        At the BEGINNING, ask the user what they need help with.
        """,
    )
)


class PolinskyTool(lr.ToolMessage):
    request: str = "polinsky"
    purpose: str = "To compute the polinsky transform of a <number>"
    number: int

    def handle(self) -> AgentDoneTool:
        p = 3 * self.number + 1
        return AgentDoneTool(content=f"The Polinsky transform of {self.number} is {p}")


class ChichikovTool(lr.ToolMessage):
    request: str = "chichikov"
    purpose: str = "To compute the Chichikov transform of a <number>"
    number: int

    def handle(self) -> AgentDoneTool:
        n = self.number**2
        return AgentDoneTool(content=f"The Chichikov transform of {self.number} is {n}")


helper = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Helper",
        system_message="""
        You have a few tools to help answer the user's questions. 
        Decide which tool to use, and send your request using the correct format 
        for the tool.
        """,
    )
)
helper.enable_message(PolinskyTool)
helper.enable_message(ChichikovTool)

main_task = lr.Task(main, interactive=False)
helper_task = lr.Task(helper, interactive=False)

main_task.add_sub_task(helper_task)

main_task.run()
</file>

<file path="examples/basic/autocorrect.py">
"""
A two agent chat system where
- AutoCorrect agent corrects the user's possibly mistyped input,
- Chatter agent responds to the corrected user's input.

Run it like this:

python3 examples/basic/autocorrect.py

"""

import typer
from rich import print

import langroid as lr
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


def chat() -> None:
    print(
        """
        [blue]Welcome to the Autocorrecting Chatbot!
        You can quickly type your message, don't even look at your keyboard. 
        Feel free to type and I will try my best to understand it,
        and I will type out what I think you meant.
        If you agree with my suggestion, just hit enter so I can respond to it.
        If you disagree with my suggestion, say "try again" or say "no" or something 
        similar, and I will try again.
        When I am confused, I will offer some numbered choices to pick from.
        
        Let's go! Enter x or q to quit at any point.
        """
    )

    config = ChatAgentConfig(
        llm=OpenAIGPTConfig(
            chat_model=OpenAIChatModel.GPT4o,
        ),
        vecdb=None,
    )
    autocorrect_agent = ChatAgent(config)
    autocorrect_task = Task(
        autocorrect_agent,
        name="AutoCorrect",
        system_message="""
        You are an expert at understanding mistyped text. You are extremely 
        intelligent, an expert in the English language, and you have common sense, 
        so no matter how badly mistyped the text is, you will know the MOST LIKELY 
        AND SENSIBLE correct version of it.
        For any text you receive, your job is to write the correct version of it, 
        and not say anything else. 
        If you are unsure, offer up to 3 numbered suggestions, and the user will pick 
        one. Once the user selects a suggestion, simply write out that version.
        Remember to ONLY suggest sensible interpretations. For example
        "Which month is the tallest in the world" is meaningless, so you should not
        ever include such a suggestion in your list.
        Start by asking me to writing something.
        """,
    )

    chat_agent = ChatAgent(config)
    chat_task = Task(
        chat_agent,
        name="Chat",
        system_message="Answer or respond very concisely, no more than 1-2 sentences!",
        done_if_no_response=[lr.Entity.LLM],
        done_if_response=[lr.Entity.LLM],
    )
    autocorrect_task.add_sub_task(chat_task)
    autocorrect_task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
            cache_type="redis",
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat-2-agent-discuss.py">
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "langroid",
# ]
# ///

"""
Give a problem statement, two agents Alice and Bob will discuss it,
and EITHER of them may return a final result via MyFinalResultTool.

Run like this (Omit model to default to GPT4o):

python3 examples/basic/chat-2-agent-discuss.py --model gemini/gemini-2.0-flash-exp

For example, try giving his problem:
What is the prime number that comes after 17?

"""

import logging

from fire import Fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.task import TaskConfig
from langroid.agent.tools.orchestration import FinalResultTool

# set info level
logging.basicConfig(level=logging.INFO)


# Any tool subclassed from FinalResultTool can be used to return the final result
# from any agent, and it will short-circuit the flow and return the result.
class MyFinalResultTool(FinalResultTool):
    request: str = "my_final_result_tool"
    purpose: str = "To present the final <result> of a discussion"
    # override this flag since it's False by default
    _allow_llm_use: bool = True

    result: str


def main(model: str = ""):
    problem = Prompt.ask(
        """
        [blue]Alice and Bob will discuss a problem.
        Please enter the problem statement:[/blue]
        """
    )

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=128_000,
        timeout=60,
    )

    logging.warning("Setting up Alice, Bob agents...")

    alice = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_config,
            name="Alice",
            system_message=f"""
            Here is a problem the user wants to solve:
            <problem>
            {problem}
            </problem>
            To solve this, you will engage in a discussion with your colleague Bob.
            At any point, if you decide the problem is solved,
            you must use the TOOL `{MyFinalResultTool.name()}` to
            return the FINAL answer to the problem. 

            In each round of the discussion, limit yourself to a CONCISE
            message.
            """,
        )
    )

    alice.enable_message(MyFinalResultTool)
    # Set `inf_loop_cycle_len` to 0, to turn OFF inf loop detection
    alice_task_config = TaskConfig(inf_loop_cycle_len=10)
    # set up alice_task to return a result of type MyFinalResultTool
    alice_task = lr.Task(alice, config=alice_task_config, interactive=False)[
        MyFinalResultTool
    ]

    bob = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_config,
            name="Bob",
            system_message=f"""
            Here is a problem the user wants to solve:
            <problem>
            {problem}
            </problem>
            To solve this, you will engage in a discussion with your colleague Alice.
            At any point, if you decide the problem is solved,
            you must use the TOOL `{MyFinalResultTool.name()}` to
            return the FINAL answer to the problem. 

            In each round of the discussion, limit yourself to a CONCISE
            message. 
            
            You will first receive a message from Alice, and you can then follow up. 
            """,
        )
    )

    bob.enable_message(MyFinalResultTool)

    bob_task = lr.Task(bob, interactive=False, single_round=True)

    # make the Con agent the sub-task of the Pro agent, so
    # they go back and forth in the arguments
    alice_task.add_sub_task(bob_task)

    result = alice_task.run("get started")

    print(
        f"""
        FINAL RESULT:
        {result.result}
        """
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/chat-azure-async-client.py">
"""
Example showing how to use Langroid with Azure OpenAI and Entra ID
authentication by providing a custom client.

This is an async version of the example in chat-azure-client.py.

For more details see here:
https://langroid.github.io/langroid/notes/custom-azure-client/
https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity

"""

import os

import azure.identity as azure_identity
import azure.identity.aio as azure_identity_async
from dotenv import load_dotenv
from openai import AsyncAzureOpenAI, AzureOpenAI

import langroid as lr
import langroid.language_models as lm

load_dotenv()


def get_azure_openai_client():
    return AzureOpenAI(
        api_version="2024-10-21",
        azure_endpoint=os.environ["AZURE_OPENAI_API_BASE"],
        azure_ad_token_provider=azure_identity.get_bearer_token_provider(
            azure_identity.DefaultAzureCredential(),
            "https://cognitiveservices.azure.com/.default",
        ),
    )


def get_azure_openai_async_client():
    return AsyncAzureOpenAI(
        api_version="2024-10-21",
        azure_endpoint=os.environ["AZURE_OPENAI_API_BASE"],
        azure_ad_token_provider=azure_identity_async.get_bearer_token_provider(
            azure_identity_async.DefaultAzureCredential(),
            "https://cognitiveservices.azure.com/.default",
        ),
    )


lm_config = lm.AzureConfig(
    azure_openai_client_provider=get_azure_openai_client,
    azure_openai_async_client_provider=get_azure_openai_async_client,
)


async def main():
    agent = lr.ChatAgent(lr.ChatAgentConfig(llm=lm_config))
    task = lr.Task(agent, interactive=False)
    response = await task.run_async(
        "Who is the president of the United States? Reply and end with DONE"
    )
    print(response)


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())
</file>

<file path="examples/basic/chat-azure-client.py">
"""
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:
https://langroid.github.io/langroid/notes/custom-azure-client/
https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity

"""

import os

from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from dotenv import load_dotenv
from openai import AzureOpenAI

import langroid as lr
import langroid.language_models as lm

load_dotenv()


def get_azure_openai_client():
    return AzureOpenAI(
        api_version="2024-10-21",
        azure_endpoint=os.environ["AZURE_OPENAI_API_BASE"],
        azure_ad_token_provider=get_bearer_token_provider(
            DefaultAzureCredential(),
            "https://cognitiveservices.azure.com/.default",
        ),
    )


lm_config = lm.AzureConfig(
    azure_openai_client_provider=get_azure_openai_client,
)

if __name__ == "__main__":
    agent = lr.ChatAgent(lr.ChatAgentConfig(llm=lm_config))
    task = lr.Task(agent, interactive=False)
    task.run("Who is the president of the United States? Reply and end with DONE")
</file>

<file path="examples/basic/chat-local-numerical.py">
"""
Test multi-round interaction with a local LLM, playing a simple "doubling game".

In each round:

- User gives a number
- LLM responds with the double of that number

Run like this --

python3 examples/basic/chat-local-numerical.py -m <local_model_name>

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import os

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.utils.configuration import settings

# for best results:
DEFAULT_LLM = lm.OpenAIChatModel.GPT4o

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# (1) Define the desired fn-call as a ToolMessage via Pydantic.


def app(
    m: str = DEFAULT_LLM,  # model name
    d: bool = False,  # debug
    nc: bool = False,  # no cache
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=4096,  # set this based on model
        max_output_tokens=100,
        temperature=0.2,
        timeout=45,
    )

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_cfg,
            system_message="""
            You are a number-doubling expert. When user gives you a NUMBER,
            simply respond with its DOUBLE and SAY NOTHING ELSE.
            DO NOT EXPLAIN YOUR ANSWER OR YOUR THOUGHT PROCESS.
            """,
        )
    )

    task = lr.Task(agent)
    task.run("15")  # initial number


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/basic/chat-local.py">
"""
Basic chat example with a local LLM.

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

Run this script as follows:

```
python examples/basic/chat-local.py --model <local_model_spec>
```

"""

from fire import Fire

import langroid as lr
import langroid.language_models as lm

# Assume you've run `ollama pull mistral` to spin up `mistral` locally.
# Notes:
# - we use `lm.OpenAIGPTConfig` to incidate this config is for LLMs served
#    at OpenAI-compatible endpoints)
# - if you omit `chat_model` below, it defaults to OpenAI GPT4-turbo,
#   or you can explicitly specify it as `lm.OpenAIChatModel.GPT4` or `lm.OpenAIChatModel.GPT4o`


def main(model: str = ""):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,  # or,e.g. "ollama/mistral"
        max_output_tokens=200,
        chat_context_length=2048,  # adjust based on your local LLM params
    )

    # Alternatively, if you've used ooba or other lib to spin up a Local LLM
    # at an OpenAI-compatible endpoint, say http://localhost:8000, you can set the
    # `chat_model` as follows (note you have to prefix it with 'local'):
    # llm_config = lm.OpenAIGPTConfig(
    #     chat_model="local/localhost:8000"
    # )
    # If the endpoint is listening at https://localhost:8000/v1, you must include the `v1`
    # at the end, e.g. chat_model="local/localhost:8000/v1"

    agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        system_message="""Be helpful but very very concise""",
    )

    agent = lr.ChatAgent(agent_config)

    task = lr.Task(agent)

    task.run()


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/chat-persist.py">
"""
Variant of chat.py, showing how you can save conversation state, end the script, and
resume the conversation later by re-running the script.

The most basic chatbot example, using the default settings.
A single Agent allows you to chat with a pre-trained Language Model.

Run like this:

python3 examples/basic/chat.py

Use optional arguments to change the settings, e.g.:

-m <local_model_spec>
-ns # no streaming
-d # debug mode
-nc # no cache
-sm <system_message>
-q <initial user msg>

For details on running with local or non-OpenAI models, see:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import logging
import pickle
from pathlib import Path

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.utils.configuration import Settings, set_global

STATE_CACHE_DIR = ".cache/agent-state"

app = typer.Typer()
logger = logging.getLogger(__name__)
# set the logging level to INFO
logger.setLevel(logging.INFO)
# Create classes for non-OpenAI model configs


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    query: str = typer.Option("", "--query", "-q", help="initial user query or msg"),
    sys_msg: str = typer.Option(
        "You are a helpful assistant. Be concise in your answers.",
        "--sysmsg",
        "-sm",
        help="system message",
    ),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to the basic chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    # use the appropriate config instance depending on model name
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=4096,
        timeout=45,
    )

    # check if history.pkl exists under STATE_CACHE_DIR, and if it does, load it
    # into agent.message_history
    hist_path = Path(STATE_CACHE_DIR) / "history.pkl"
    hist_found = False
    try:
        if hist_path.exists():
            # read the history from the cache
            with open(str(hist_path), "rb") as f:
                msg_history = pickle.load(f)
            n_msgs = len(msg_history)
            logger.info(f"Loaded {n_msgs} messages from cache")
            hist_found = True
        else:
            sys_msg = Prompt.ask(
                "[blue]Tell me who I am. Hit Enter for default, or type your own\n",
                default=sys_msg,
            )

    except Exception:
        logger.warning("Failed to load message history from cache")
        pass

    config = ChatAgentConfig(
        system_message=sys_msg,
        llm=llm_config,
    )
    agent = ChatAgent(config)

    if hist_found:
        # overrides sys_msg set in config
        agent.message_history = msg_history

    # use restart=False so the state is not cleared out at start,
    # which allows continuing the conversation.
    task = Task(agent, restart=False)
    # OpenAI models are ok with just a system msg,
    # but in some scenarios, other (e.g. llama) models
    # seem to do better when kicked off with a sys msg and a user msg.
    # In those cases we may want to do task.run("hello") instead.
    if query:
        task.run(query)
    else:
        task.run()

    # Create STATE_CACHE_DIR if it doesn't exist
    Path(STATE_CACHE_DIR).mkdir(parents=True, exist_ok=True)
    # Save the conversation state to hist_path
    with open(str(hist_path), "wb") as f:
        pickle.dump(agent.message_history, f)
    logger.info(f"Saved {len(agent.message_history)} messages to cache")


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat-search-assistant-local.py">
"""
Version of chat-search-assistant.py that uses local LLMs.
Tested and works ok nous-hermes2-mixtral, but still has issues.

3-Agent system where:
- Assistant takes user's (complex) question, breaks it down into smaller pieces
    if needed
- Searcher takes Assistant's question, uses the Search tool to search the web
    (using DuckDuckGo), and returns a coherent answer to the Assistant.
- Critic takes Assistant's final answer, and provides feedback on it.

Once the Assistant thinks it has enough info to answer the user's question, it
says DONE and presents the answer to the user.

See also: chat-search for a basic single-agent search

python3 examples/basic/chat-search-assistant.py

There are optional args, especially note these:

-m <model_name>: to run with a different LLM model (default: gpt4-turbo)

For example try this question:

during which years did Beethoven live, and does his life overlap with that of Liszt?

You can specify a local in a few different ways, e.g. `-m local/localhost:8000/v1`
or `-m ollama/mistral` etc. See here how to use Langroid with local LLMs:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


"""

from typing import List, Optional, Type

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.duckduckgo_search_tool import DuckduckgoSearchTool
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


class QuestionTool(lr.ToolMessage):
    request: str = "question_tool"
    purpose: str = "Ask a SINGLE <question> that can be answered from a web search."
    question: str

    @classmethod
    def examples(cls) -> List[lr.ToolMessage]:
        return [
            cls(question="Which superconductor material was discovered in 2023?"),
            cls(question="What AI innovation did Meta achieve in 2024?"),
        ]


class FinalAnswerTool(lr.ToolMessage):
    request: str = "final_answer_tool"
    purpose: str = """
        Present the intermediate <steps> and 
        final <answer> to the user's original query.
        """
    steps: str
    answer: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            cls(
                steps="1. Man is mortal. 2. Plato was a man.",
                answer="Plato was mortal.",
            ),
            cls(
                steps="1. The moon landing was in 1969. 2. Kennedy was president "
                "during 1969.",
                answer="Kennedy was president during the moon landing.",
            ),
        ]


class FeedbackTool(lr.ToolMessage):
    request: str = "feedback_tool"
    purpose: str = "Provide <feedback> on the user's answer."
    feedback: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            cls(feedback=""),
            cls(
                feedback="""
                The answer is invalid because the conclusion does not follow from the
                steps. Please check your reasoning and try again.
                """
            ),
        ]


class AssistantAgent(lr.ChatAgent):
    n_questions: int = 0  # how many questions in THIS round
    has_asked: bool = False  # has ANY question been asked
    original_query: str | None = None

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.USER:
            # either first query from user, or returned result from Searcher
            self.n_questions = 0  # reset search count

        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            if self.has_asked:
                return f"""
                You may have intended to use a tool, but your JSON format may be wrong.
                
                REMINDER: You must do one of the following:
                - If you are ready with the final answer to the user's ORIGINAL QUERY
                    [ Remember it was: {self.original_query} ],
                  then present your reasoning steps and final answer using the 
                  `final_answer_tool` in the specified JSON format.
                - If you still need to ask a question, then use the `question_tool`
                  to ask a SINGLE question that can be answered from a web search.
                """
            elif self.original_query is not None:
                return f"""
                You must ask a question using the `question_tool` in the specified format,
                to break down the user's original query: {self.original_query} into 
                smaller questions that can be answered from a web search.
                """

    def question_tool(self, msg: QuestionTool) -> str:
        self.n_questions += 1
        self.has_asked = True
        if self.n_questions > 1:
            # there was already a search, so ignore this one
            return ""
        # valid question tool: re-create it so Searcher gets it
        return msg.to_json()

    def final_answer_tool(self, msg: FinalAnswerTool) -> str:
        if not self.has_asked or self.n_questions > 1:
            # not yet asked any questions, or LLM is currently asking
            # a question (and this is the second one in this turn, and so should
            # be ignored), ==>
            # cannot present final answer yet (LLM may have hallucinated this json)
            return ""
        # valid final answer tool: PASS it on so Critic gets it
        return lr.utils.constants.PASS_TO + "Critic"

    def feedback_tool(self, msg: FeedbackTool) -> str:
        if msg.feedback == "":
            return lr.utils.constants.DONE
        else:
            return f"""
            Below is feedback about your answer. Take it into account to 
            improve your answer, and present it again using the `final_answer_tool`.
            
            FEEDBACK:
            
            {msg.feedback}
            """

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if self.original_query is None:
            self.original_query = (
                message if isinstance(message, str) else message.content
            )
        result = super().llm_response(message)
        if result is None:
            return result
        # result.content may contain a premature DONE
        # (because weak LLMs tend to repeat their instructions)
        # We deem a DONE to be accidental if no search query results were received
        if not isinstance(message, ChatDocument) or not (
            message.metadata.sender_name == "Searcher"
        ):
            # no search results received yet, so should NOT say DONE
            if isinstance(result, str):
                return result.content.replace(lr.utils.constants.DONE, "")
            result.content = result.content.replace(lr.utils.constants.DONE, "")
            return result

        return result


class CriticAgent(lr.ChatAgent):
    def final_answer_tool(self, msg: FinalAnswerTool) -> str:
        # received from Assistant. Extract the components as plain text,
        # so that the Critic LLM can provide feedback
        return f"""
        The user has presented the following intermediate steps and final answer
        shown below. Please provide feedback using the `feedback_tool`.
        Remember to set the `feedback` field to an empty string if the answer is valid,
        otherwise give specific feedback on what the issues are and how the answer 
        can be improved.
        
        STEPS: {msg.steps}
        
        ANSWER: {msg.answer}
        """

    def feedback_tool(self, msg: FeedbackTool) -> str:
        # say DONE and PASS to the feedback goes back to Assistant to handle
        return lr.utils.constants.DONE + " " + lr.utils.constants.PASS


class SearcherAgentConfig(lr.ChatAgentConfig):
    search_tool_class: Type[lr.ToolMessage]


class SearcherAgent(lr.ChatAgent):
    n_searches: int = 0
    curr_query: str | None = None

    def __init__(self, config: SearcherAgentConfig):
        super().__init__(config)
        self.config: SearcherAgentConfig = config
        self.enable_message(config.search_tool_class)
        self.enable_message(QuestionTool, use=False, handle=True)

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if (
            isinstance(msg, ChatDocument)
            and msg.metadata.sender == lr.Entity.LLM
            and self.n_searches == 0
        ):
            search_tool_name = self.config.search_tool_class.default_value("request")
            return f"""
            You forgot to use the web search tool to answer the 
            user's question : {self.curr_query}.
            REMEMBER - you must ONLY answer the user's questions based on 
             results from a web-search, and you MUST NOT ANSWER them yourself.
             
            Please use the `{search_tool_name}` tool 
            using the specified JSON format, then compose your answer.
            """

    def question_tool(self, msg: QuestionTool) -> str:
        self.curr_query = msg.question
        search_tool_name = self.config.search_tool_class.default_value("request")
        return f"""
        User asked this question: {msg.question}.
        Perform a web search using the `{search_tool_name}` tool
        using the specified JSON format, to find the answer.
        """

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if (
            isinstance(message, ChatDocument)
            and message.metadata.sender == lr.Entity.AGENT
            and self.n_searches > 0
        ):
            # must be search results from the web search tool,
            # so let the LLM compose a response based on the search results
            self.n_searches = 0  # reset search count

            result = super().llm_response_forget(message)
            # Augment the LLM's composed answer with a helpful nudge
            # back to the Assistant
            result.content = f"""
            Here are the web-search results for the question: {self.curr_query}.
            ===
            {result.content}
            ===
            Decide if you want to ask any further questions, for the 
            user's original question.             
            """
            self.curr_query = None
            return result

        # Handling query from user (or other agent)
        result = super().llm_response_forget(message)
        if result is None:
            return result
        tools = self.get_tool_messages(result)
        if all(not isinstance(t, self.config.search_tool_class) for t in tools):
            # LLM did not use search tool;
            # Replace its response with a placeholder message
            # and the agent fallback_handler will remind the LLM
            result.content = "Did not use web-search tool."
            return result

        self.n_searches += 1
        # result includes a search tool, but may contain DONE in content,
        # so remove that
        result.content = result.content.replace(lr.utils.constants.DONE, "")
        return result


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    print(
        """
        [blue]Welcome to the Web Search Assistant chatbot!
        I will try to answer your complex questions. 
        
        Enter x or q to quit at any point.
        """
    )
    load_dotenv()

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )

    assistant_config = lr.ChatAgentConfig(
        system_message="""
        You are a resourceful assistant, able to think step by step to answer
        complex questions from the user. You must break down complex questions into
        simpler questions that can be answered by a web search. You must ask me 
        (the user) each question ONE BY ONE, using the `question_tool` in
         the specified format, and I will do a web search and send you
        a brief answer. Once you have enough information to answer my original
        (complex) question, you MUST present your INTERMEDIATE STEPS and FINAL ANSWER
        using the `final_answer_tool` in the specified JSON format.
        You will then receive FEEDBACK from the Critic, and if needed
        you should try to improve your answer based on this feedback.
        """,
        llm=llm_config,
        vecdb=None,
    )
    assistant_agent = AssistantAgent(assistant_config)
    assistant_agent.enable_message(QuestionTool)
    assistant_agent.enable_message(FinalAnswerTool)
    assistant_agent.enable_message(FeedbackTool, use=False, handle=True)

    search_tool_handler_method = DuckduckgoSearchTool.default_value("request")

    search_agent_config = SearcherAgentConfig(
        search_tool_class=DuckduckgoSearchTool,
        llm=llm_config,
        vecdb=None,
        system_message=f"""
        You are a web-searcher. For ANY question you get, you must use the
        `{search_tool_handler_method}` tool/function-call to get up to 5 results.
        Once you receive the results, you must compose a CONCISE answer 
        based on the search results and say DONE and show the answer to me,
        along with references, in this format:
        DONE [... your CONCISE answer here ...]
        SOURCES: [links from the web-search that you used]
        
        EXTREMELY IMPORTANT: DO NOT MAKE UP ANSWERS, ONLY use the web-search results.
        """,
    )
    search_agent = SearcherAgent(search_agent_config)

    assistant_task = lr.Task(
        assistant_agent,
        name="Assistant",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    search_task = lr.Task(
        search_agent,
        name="Searcher",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )

    critic_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message="""
        You excel at logical reasoning and combining pieces of information.
        The user will send you a summary of the intermediate steps and final answer.
        You must examine these and provide feedback to the user, using the 
        `feedback_tool`, as follows:
        - If you think the answer is valid, 
            simply set the `feedback` field to an empty string "".
        - Otherwise set the `feedback` field to a reason why the answer is invalid,
            and suggest how the user can improve the answer.
        """,
    )
    critic_agent = CriticAgent(critic_agent_config)
    critic_agent.enable_message(FeedbackTool)
    critic_agent.enable_message(FinalAnswerTool, use=False, handle=True)
    critic_task = lr.Task(
        critic_agent,
        name="Critic",
        interactive=False,
    )
    assistant_task.add_sub_task([search_task, critic_task])
    question = Prompt.ask("What do you want to know?")
    assistant_task.run(question)


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat-search-assistant.py">
"""
2-Agent system where:
- Assistant takes user's (complex) question, breaks it down into smaller pieces
    if needed
- WebSearcher takes Assistant's question, uses the Search tool to search the web
    (default DuckDuckGo, or Google or Metaphor as specified by user), and returns a
    coherent answer to the Assistant.

Once the Assistant thinks it has enough info to answer the user's question, it
says DONE and presents the answer to the user.

See also: chat-search for a basic single-agent search

python3 examples/basic/chat-search-assistant.py

There are optional args, especially note these:

-p or --provider: google or ddg or metaphor (default: google)
-m <model_name>: to run with a different LLM model (default: gpt4-turbo)

You can specify a local in a few different ways, e.g. `-m local/localhost:8000/v1`
or `-m ollama/mistral` etc. See here how to use Langroid with local LLMs:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


NOTE:
(a) If using Google Search, you must have GOOGLE_API_KEY and GOOGLE_CSE_ID
environment variables in your `.env` file, as explained in the
[README](https://github.com/langroid/langroid#gear-installation-and-setup).

(b) If using MetaphorSearchTool, you need to:
* set the METAPHOR_API_KEY environment variables in
your `.env` file, e.g. `METAPHOR_API_KEY=your_api_key_here`
* install langroid with the `metaphor` extra, e.g.
`pip install langroid[metaphor]` or `uv pip install langroid[metaphor]`
`poetry add langroid[metaphor]` or `uv add langroid[metaphor]`
(it installs the `metaphor-python` package from pypi).
For more information, please refer to the official docs:
https://metaphor.systems/

"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.duckduckgo_search_tool import DuckduckgoSearchTool
from langroid.agent.tools.google_search_tool import GoogleSearchTool
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    provider: str = typer.Option(
        "ddg",
        "--provider",
        "-p",
        help="search provider name (google, metaphor, ddg)",
    ),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    print(
        """
        [blue]Welcome to the Web Search Assistant chatbot!
        I will try to answer your complex questions. 
        
        Enter x or q to quit at any point.
        """
    )
    load_dotenv()

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=8_000,
        temperature=0,
        max_output_tokens=200,
        timeout=45,
    )

    assistant_config = lr.ChatAgentConfig(
        system_message=f"""
        You are a resourceful assistant, able to think step by step to answer
        complex questions from the user. You must break down complex questions into
        simpler questions that can be answered by a web search. You must ask me 
        (the user) each question ONE BY ONE, and I will do a web search and send you
        a brief answer. Once you have enough information to answer my original
        (complex) question, you MUST say {DONE} and present the answer to me.
        """,
        llm=llm_config,
        vecdb=None,
    )
    assistant_agent = lr.ChatAgent(assistant_config)

    match provider:
        case "google":
            search_tool_class = GoogleSearchTool
        case "metaphor":
            from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool

            search_tool_class = MetaphorSearchTool
        case "ddg":
            search_tool_class = DuckduckgoSearchTool
        case _:
            raise ValueError(f"Unsupported provider {provider} specified.")

    search_tool_handler_method = search_tool_class.name()

    search_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message=f"""
        You are a web-searcher. For any question you get, you must use the TOOL
        `{search_tool_handler_method}`  to get up to 5 results.
        I WILL SEND YOU THE RESULTS; DO NOT MAKE UP THE RESULTS!!
        Once you receive the results, you must compose a CONCISE answer 
        based on the search results and say {DONE} and show the answer to me,
        in this format:
        {DONE} [... your CONCISE answer here ...]
        IMPORTANT:
        * YOU MUST WAIT FOR ME TO SEND YOU THE SEARCH RESULTS BEFORE saying  {DONE}.
        * YOU Can only use the TOOL `{search_tool_handler_method}` 
            ONE AT A TIME, even if you get multiple questions!
        """,
    )
    search_agent = lr.ChatAgent(search_agent_config)
    search_agent.enable_message(search_tool_class)

    assistant_task = lr.Task(
        assistant_agent,
        name="Assistant",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    search_task = lr.Task(
        search_agent,
        name="Searcher",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    assistant_task.add_sub_task(search_task)
    question = Prompt.ask("What do you want to know?")
    assistant_task.run(question)


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat-tool-function.py">
"""
Bare bones example of using tool/function-call

Run like this, optionally specifying an LLM:

python3 examples/basic/chat-tool-function.py

or

python3 examples/basic/chat-tool-function.py -m ollama/mistral:7b-instruct-v0.2-q8_0

or 

uv run examples/basic/chat-tool-function.py -m deepseek/deepseek-reasoner

"""

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import FinalResultTool
from pydantic import BaseModel, Field

# define a nested structure for Company information


class CompanyFinancials(BaseModel):
    shares: int = Field(..., description="shares outstanding of company")
    price: float = Field(..., description="price per share of company")
    eps: float = Field(..., description="earnings per share of company")


class CompanyInfo(BaseModel):
    name: str = Field(..., description="name of company")
    industry: str = Field(..., description="industry of company")
    financials: CompanyFinancials = Field(..., description="financials of company")


# define a ToolMessage corresponding to the above structure


class CompanyInfoTool(lr.agent.ToolMessage):
    request: str = "company_info_tool"  # agent method that handles this tool
    purpose: str = (
        "To extract <company_info> from a passage and compute market-capitalization."
    )
    company_info: CompanyInfo

    @classmethod
    def examples(cls):
        """Examples that will be compiled to few-shot examples for the LLM.
        Illustrating two types of examples below:
        - example instance
        - (description, example) tuple
        """
        return [
            cls(
                company_info=CompanyInfo(
                    name="IBM",
                    industry="Technology",
                    financials=CompanyFinancials(shares=1.24e9, price=140.15, eps=4.68),
                )
            ),
            (
                "I want to extract and present company info from the passage",
                cls(
                    company_info=CompanyInfo(
                        name="Apple",
                        industry="Technology",
                        financials=CompanyFinancials(
                            shares=16.82e9, price=149.15, eps=5.68
                        ),
                    )
                ),
            ),
        ]

    def handle(self) -> FinalResultTool:
        """Handle LLM's structured output if it matches CompanyInfo structure.
        This suffices for a "stateless" tool.
        If the tool handling requires agent state, then
        instead of this `handle` method, define a `company_info_tool`
        method in the agent.
        """
        mkt_cap = (
            self.company_info.financials.shares * self.company_info.financials.price
        )
        print(
            f"""
            Got Valid Company Info.
            The market cap of {self.company_info.name} is ${mkt_cap/1e9}B.
            """
        )
        return FinalResultTool(
            market_cap=mkt_cap,
            info=self.company_info,
        )


def run(model: str = ""):  # or, e.g., "ollama/mistral:7b-instruct-v0.2-q8_0"
    lm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,  # or
    )
    tool_name = CompanyInfoTool.default_value("request")
    agent_config = lr.ChatAgentConfig(
        llm=lm_config,
        system_message=f"""
        You are a company-info extraction expert. When user gives you a TEXT PASSAGE,
        simply extract the company information and 
        present it using the `{tool_name}` tool/function-call.
        """,
    )
    agent = lr.ChatAgent(agent_config)
    agent.enable_message(CompanyInfoTool)

    # text to present to the LLM
    paragraph = """
        Apple Inc. is an American multinational technology company that specializes in 
        consumer electronics, computer software, and online services.
        It has shares outstanding of 16.82 billion, and a price per share of $149.15.
        The earnings per share is $5.68.
        """

    # test 1:
    # see that the LLM extracts the company information and presents it using the tool
    response = agent.llm_response(paragraph)

    print(response.content)

    # test 2:
    # wrap the agent in a Task, so that the ToolMessage is handled by the handle method
    task = lr.Task(agent, interactive=False)
    result = task[FinalResultTool].run(paragraph)
    assert result.market_cap > 0
    assert "Apple" in result.info.name


if __name__ == "__main__":
    Fire(run)
</file>

<file path="examples/basic/chat-tree-structured-simple.py">
"""
Simple example showing tree-structured computation, 
a variation of `examples/basic/chat-tree.py` which uses strict output formatting
to reliably wrap calls to agents in standard Python functions, allowing
explicit control over control flow.

The task consists of performing this calculation for a given input number n:

def Main(n):
    if n is odd:
        return (3*n+1) + n
    else:
        If n is divisible by 10:
            return n/10 + n
        else:
            return n/2 + n

Each step is performed by an LLM call, and strict output formatting ensures that
a valid typed response is returned (rather than a string which requires another
LLM call to interpret).

We evaluate the conditions with a `condition_agent` which is given an integer and
a condition and return a Boolean and evaluate the transformations of `n` with
a `transformation_agent` which is given an integer and a transformation rule
and returns the transformed integer.

Finally, we add the result with the original `n` using an `adder_agent` which
illustrates strict output usage in `Task`s.

For more details on structured outputs, see the notes at
https://langroid.github.io/langroid/notes/structured-output/.
"""

import typer
from rich.prompt import Prompt

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


def chat() -> int:
    condition_agent = ChatAgent(
        ChatAgentConfig(
            system_message="""
            You will be provided with a condition and a
            number; your goal is to determine whether
            that number satisfies the condition.

            Respond in JSON format, with `value` set
            to the result.
            """,
            output_format=bool,
        )
    )
    transformation_agent = ChatAgent(
        ChatAgentConfig(
            system_message="""
            You will be provided with a number and an
            transformation of the number to perform.

            Respond in JSON format, with `value` set
            to the result.
            """,
            output_format=int,
        )
    )

    def check_condition(n: int, condition: str) -> bool:
        output = condition_agent.llm_response_forget(
            f"""
            Number: {n}
            Condition: {condition}
            """
        )
        return condition_agent.from_ChatDocument(output, bool)  # type: ignore

    def apply_transformation(n: int, transformation: str) -> int:
        output = transformation_agent.llm_response_forget(
            f"""
            Number: {n}
            Transformation: {transformation}
            """
        )
        return transformation_agent.from_ChatDocument(output, int)  # type: ignore

    num = int(Prompt.ask("Enter a number"))
    is_even = check_condition(num, "The number is even.")

    if is_even:
        is_divisible_by_10 = check_condition(num, "The number is divisible by 10.")

        if is_divisible_by_10:
            to_adder = apply_transformation(num, "n/10 where the number is n.")
        else:
            to_adder = apply_transformation(num, "n/2 where the number is n.")
    else:
        to_adder = apply_transformation(num, "3n+1 where the number is n.")

    class AddNumTool(ToolMessage):
        request: str = "add_num"
        purpose: str = "Add <number> to the original number, return the result"
        number: int

        def handle(self) -> str:
            total = num + self.number
            return f"{DONE} {total}"

    # We could also have the agent output a the call in a single step and handle
    # it ourselves (or apply it immediately)
    adder_agent = ChatAgent(
        ChatAgentConfig(
            system_message="""
            You will be given a number n.
            You have to add it to the original number and return the result.
            You do not know the original number, so you must use the 
            `add_num` tool/function for this. 
            """,
            output_format=AddNumTool,
        )
    )
    adder_agent.enable_message(AddNumTool)
    adder_task = Task(adder_agent, interactive=False, name="Adder")

    # compute the final output value
    return adder_task[int].run(str(to_adder))  # type: ignore


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat-tree-structured.py">
"""
Simple example showing tree-structured computation
where each node in the tree is handled by a separate agent.
A variation of `examples/basic/chat-tree.py` which uses strict output formatting
and agent logic to enforce the behavior specified in the prompts.

See the use of `set_output_format()` in ConditionalAgent.

The task consists of performing this calculation for a given input number n:

def Main(n):
    if n is odd:
        return (3*n+1) + n
    else:
        If n is divisible by 10:
            return n/10 + n
        else:
            return n/2 + n

To make this "interesting", we represent this computation hierarchically,
in the form of this tree:

Main
- Odd
    - Adder
- Even
    - EvenZ
        - Adder
    - EvenNZ
        - Adder

For a full write-up on the design considerations, see the documentation page on
Hiearchical Agent Computations at https://langroid.github.io/langroid/examples/agent-tree/

For more details on structured outputs, see the notes at
https://langroid.github.io/langroid/notes/structured-output/.
"""

import typer
from rich.prompt import Prompt

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.chat_document import ChatDocument
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE
from langroid.utils.globals import GlobalState
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


class MyGlobalState(GlobalState):
    number: int | None = None


class AskNumTool(ToolMessage):
    request: str = "ask_num"
    purpose: str = "Ask user for the initial number"

    def handle(self) -> str:
        """
        This is a stateless tool (i.e. does not use any Agent member vars), so we can
        define the handler right here, instead of defining an `ask_num`
        method in the agent.
        """
        num = int(Prompt.ask("Enter a number"))
        # record this in global state, so other agents can access it
        MyGlobalState.set_values(number=num)
        return str(num)


class AddNumTool(ToolMessage):
    request: str = "add_num"
    purpose: str = "Add <number> to the original number, return the result"
    number: int

    def handle(self) -> AgentDoneTool:
        """
        This is a stateless tool (i.e. does not use any Agent member vars), so we can
        define the handler right here, instead of defining an `add_num`
        method in the agent.
        """
        total = MyGlobalState.get_value("number") + self.number
        return AgentDoneTool(
            tools=[ResultTool(result=total)],
        )


class MatchTool(ToolMessage):
    request: str = "match"
    purpose: str = "To express whether the input number matches your condition."
    matches: bool


class ResultTool(ToolMessage):
    request: str = "result"
    purpose: str = (
        "To express the result of your transformation applied to the input number."
    )
    result: int


class ConditionalAgentConfig(ChatAgentConfig):
    top_level: bool = False


class ConditionalAgent(ChatAgent):
    def __init__(self, config: ConditionalAgentConfig = ConditionalAgentConfig()):
        super().__init__(config)
        self.config: ConditionalAgentConfig = config  # type: ignore
        # Should the next request be treated as self-generated?
        self.generated_request: bool = False

        if self.config.top_level:
            # We always begin by requesting a number from the user
            self.set_output_format(AskNumTool)
            self.enable_message(AskNumTool)
            self.enable_message(ResultTool, handle=True, use=False)
        else:
            self.enable_message([MatchTool, ResultTool])
            # We always begin by checking whether the number matches the agent's condiditon
            self.set_output_format(MatchTool)

    def ask_num(self, msg: AskNumTool) -> str:
        self.set_output_format(None)
        return msg.handle()

    def match(self, msg: MatchTool) -> str:
        if not msg.matches:
            return DONE

        # The agent must next return the transformed number
        self.set_output_format(ResultTool)
        self.generated_request = True
        return "Now, return the input number, after applying your transformation."

    def result(self, msg: ResultTool) -> str | ChatDocument | AgentDoneTool:
        if self.config.top_level:
            self.set_output_format(AskNumTool)
            # Return the answer if we are the top-level task
            return f"{DONE} {msg.result}"
        elif self.generated_request:
            self.generated_request = False
            return self.create_llm_response(
                content=str(msg.result),
            )
        else:
            self.set_output_format(MatchTool)

        # Propogate the result up if we are done
        return AgentDoneTool(
            tools=[msg],
        )


def chat() -> None:
    main_task = Task(
        ConditionalAgent(
            ConditionalAgentConfig(
                top_level=True,
            )
        ),
        interactive=False,
        name="Main",
        system_message="""
        You will ask the user for a number with the `ask_num` tool; you should respond with exactly that number,
        say nothing else.
        """,
    )

    prompt_format = """
        You will receive a number; you should first check whether that number
        matches your condition.

        Condition: {condition}

        If so, you should respond with a transformed version of the number:

        Transformation: {transformation}
        """

    even_task = Task(
        ConditionalAgent(),
        interactive=False,
        name="Even",
        system_message=prompt_format.format(
            condition="The number is even.",
            transformation="Nothing, return the number you were provided.",
        ),
    )
    evenz_task = Task(
        ConditionalAgent(),
        interactive=False,
        name="EvenZ",
        system_message=prompt_format.format(
            condition="The number is divisible by 10.",
            transformation="Return n/10 where n is the provided number.",
        ),
    )
    even_nz_task = Task(
        ConditionalAgent(),
        interactive=False,
        name="EvenNZ",
        system_message=prompt_format.format(
            condition="The number is not divisible by 10.",
            transformation="Return n/2 where n is the provided number.",
        ),
    )
    odd_task = Task(
        ConditionalAgent(),
        interactive=False,
        name="Odd",
        system_message=prompt_format.format(
            condition="The number is odd.",
            transformation="Return n*3 + 1",
        ),
    )

    adder_agent = ChatAgent()
    adder_agent.enable_message(AddNumTool)
    adder_task = Task(
        # ensure that the agent calls the tool:
        # agent[T] is a copy of agent which always outputs values of type T
        adder_agent[AddNumTool],
        name="Adder",
        interactive=False,
        system_message="""
        You will be given a number n.
        You have to add it to the original number and return the result.
        You do not know the original number, so you must use the 
        `add_num` tool/function for this. 
        """,
    )

    # set up tasks and subtasks
    main_task.add_sub_task([even_task, odd_task])
    even_task.add_sub_task([evenz_task, even_nz_task])
    evenz_task.add_sub_task(adder_task)
    even_nz_task.add_sub_task(adder_task)
    odd_task.add_sub_task(adder_task)

    # start the chat
    main_task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat-tree.py">
"""
Simple example showing tree-structured computation
where each node in the tree is handled by a separate agent.

This task consists of performing this calculation for a given input number n:

def Main(n):
    if n is odd:
        return (3*n+1) + n
    else:
        If n is divisible by 10:
            return n/10 + n
        else:
            return n/2 + n

To make this "interesting", we represent this computation hierarchically,
in the form of this tree:

Main
- Odd
    - Adder
- Even
    - EvenZ
        - Adder
    - EvenNZ
        - Adder

For a full write-up on the design considerations, see the documentation page on
Hiearchical Agent Computations at https://langroid.github.io/langroid/examples/agent-tree/
"""

import typer
from rich.prompt import Prompt

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE
from langroid.utils.globals import GlobalState
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


class MyGlobalState(GlobalState):
    number: int | None = None


class AskNumTool(ToolMessage):
    request: str = "ask_num"
    purpose: str = "Ask user for the initial number"

    def handle(self) -> str:
        """
        This is a stateless tool (i.e. does not use any Agent member vars), so we can
        define the handler right here, instead of defining an `ask_num`
        method in the agent.
        """
        num = Prompt.ask("Enter a number")
        # record this in global state, so other agents can access it
        MyGlobalState.set_values(number=num)
        return str(num)


class AddNumTool(ToolMessage):
    request: str = "add_num"
    purpose: str = "Add <number> to the original number, return the result"
    number: int

    def handle(self) -> str:
        """
        This is a stateless tool (i.e. does not use any Agent member vars), so we can
        define the handler right here, instead of defining an `add_num`
        method in the agent.
        """
        return str(int(MyGlobalState.get_value("number")) + int(self.number))


def chat(model: str = "") -> None:
    config = ChatAgentConfig(
        llm=OpenAIGPTConfig(
            chat_model=model or OpenAIChatModel.GPT4o,
        ),
        vecdb=None,
    )

    main_agent = ChatAgent(config)
    main_task = Task(
        main_agent,
        name="Main",
        interactive=False,
        system_message="""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        In this case simply write the <number>, say nothing else.
        
        RESULT Message format: RESULT <number>
        In this case simply say "DONE <number>", e.g.:
        DONE 19

        To start off, ask the user for the initial number, 
        using the `ask_num` tool/function.
        """,
    )

    # Handles only even numbers
    even_agent = ChatAgent(config)
    even_task = Task(
        even_agent,
        name="Even",
        interactive=False,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if the <number> is odd, say '{DONE}'
        - otherwise, simply write the <number>, say nothing else.
        
        RESULT Message format: RESULT <number>
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    # handles only even numbers ending in Zero
    evenz_agent = ChatAgent(config)
    evenz_task = Task(
        evenz_agent,
        name="EvenZ",
        interactive=False,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if <number> n is even AND divisible by 10, compute n/10 and pass it on,
        - otherwise, say '{DONE}'
        
        RESULT Message format: RESULT <number>
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    # Handles only even numbers NOT ending in Zero
    even_nz_agent = ChatAgent(config)
    even_nz_task = Task(
        even_nz_agent,
        name="EvenNZ",
        interactive=False,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if <number> n is even AND NOT divisible by 10, compute n/2 and pass it on,
        - otherwise, say '{DONE}'
        
        RESULT Message format: RESULT <number>
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    # Handles only odd numbers
    odd_agent = ChatAgent(config)
    odd_task = Task(
        odd_agent,
        name="Odd",
        interactive=False,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if <number> n is odd, compute n*3+1 and write it.
        - otherwise, say '{DONE}'

        RESULT Message format: RESULT <number>        
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    adder_agent = ChatAgent(config)
    adder_task = Task(
        adder_agent,
        name="Adder",
        interactive=False,
        system_message="""
        You will be given a number n.
        You have to add it to the original number and return the result.
        You do not know the original number, so you must use the 
        `add_num` tool/function for this. 
        When you receive the result, say "DONE RESULT <result>", e.g.
        DONE RESULT 19
        """,
    )

    # set up tasks and subtasks
    main_task.add_sub_task([even_task, odd_task])
    even_task.add_sub_task([evenz_task, even_nz_task])
    evenz_task.add_sub_task(adder_task)
    even_nz_task.add_sub_task(adder_task)
    odd_task.add_sub_task(adder_task)

    # set up the tools
    main_agent.enable_message(AskNumTool)
    adder_agent.enable_message(AddNumTool)

    # start the chat
    main_task.run()


@app.command()
def main(
    model: str = typer.Option("", "--model", "-m", help="model to use"),
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat(model)


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/chat.py">
# script
# requires-python = ">=3.11"
# dependencies = [
#     "langroid",
# ]
# ///
"""
The most basic chatbot example, using the default settings.
A single Agent allows you to chat with a pre-trained Language Model.

Run like this:

python3 examples/basic/chat.py

Use optional arguments to change the settings, e.g.:

-m <local_model_spec>
-ns # no streaming
-d # debug mode
-nc # no cache
-sm <system_message>
-q <initial user msg>

For details on running with local or non-OpenAI models, see:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()

# Create classes for non-OpenAI model configs


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    sys_msg: str = typer.Option(
        "You are a helpful assistant. Be concise in your answers.",
        "--sysmsg",
        "-sm",
        help="system message",
    ),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to the basic chatbot!
        Enter x or q to quit at any point.
        """
    )

    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=16_000,  # set based on model
        timeout=45,
    )

    sys_msg = Prompt.ask(
        "[blue]Tell me who I am. Hit Enter for default, or type your own\n",
        default=sys_msg,
    )

    config = ChatAgentConfig(
        system_message=sys_msg,
        llm=llm_config,
    )
    agent = ChatAgent(config)
    task = Task(agent)
    task.run("hello")


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/completion.py">
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "langroid",
# ]
# ///
"""
Interact with a base completion model, specifically the original GPT-3 base model
(i.e. davinci-002 or babbage-002),
one that has not been instruct-tuned for chat-like conversation.
This uses the legacy OpenAI Completion API.
This API simply takes pure text (NOT dialog) , and returns the LLM's completion.
Note there is no notion of system message here.

Run like this:

python3 examples/basic/completion.py

Use optional arguments to change the settings, e.g.:

-m <local_model_spec>
-ns # no streaming
-d # debug mode
-nc # no cache


For details on running with local or non-OpenAI models, see:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


def multiline_input(prompt_text):
    lines = []
    while True:
        line = Prompt.ask(prompt_text)
        if not line:
            break
        lines.append(line)
    return "\n".join(lines)


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to the basic completion engine.
        Text you enter will be completed by an LLM 
        (Default is a GPT3-class LLM, davinci-002). 
        You can enter multi-line inputs; Enter return TWICE to send your message.
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    # use the appropriate config instance depending on model name
    llm_config = lm.OpenAIGPTConfig(
        completion_model=model or "davinci-002",  # or "babbage-002"
        chat_context_length=4096,
        timeout=45,
        use_chat_for_completion=False,
    )
    llm = lm.OpenAIGPT(llm_config)

    print()
    while True:
        print("\n")
        user_msg = multiline_input("[blue]You[/blue]")
        if user_msg.lower() in ["q", "x"]:
            break
        print("\nBot: ")
        response = llm.generate(prompt=user_msg, max_tokens=50)

        if response.cached:
            print(f"[red](Cached)[/red] [green] {response.message}[/green]")


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/concurrent-tasks.py">
"""
Toy example showing how to combine results from multiple tasks running concurrently.

- main agent/task uses `multi_task_tool` tool to specify what to send to tasks t2, t3
- t2, t3 are run concurrently
- results from t2, t3 are combined and returned to main agent/task
- main agent/task then uses the combined results to generate a final response
"""

from typing import Dict

from fire import Fire
import langroid as lr
import langroid.language_models as lm
from langroid.agent.batch import run_batch_task_gen
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.globals import GlobalState

CITY_AGENT_NAME = "CityAgent"
NAME_AGENT_NAME = "NameAgent"


class MyGlobalState(GlobalState):
    name_task_map: Dict[str, str] = {}


class MultiTaskTool(lr.ToolMessage):
    request: str = "multi_task_tool"
    purpose: str = """
        Specify messages to send to multiple agents, via <agent_msgs>
        which is a dict mapping agent names to messages.
    """
    agent_msgs: Dict[str, str]

    def handle(self) -> AgentDoneTool:
        inputs = list(self.agent_msgs.values())
        agent_names = list(self.agent_msgs.keys())
        name_task_map = MyGlobalState.get_value("name_task_map")
        tasks = [name_task_map[name] for name in agent_names]

        def result2content_fn(chat_doc: lr.ChatDocument) -> str:
            return chat_doc.content

        def task_gen(i: int):  # task generator
            return tasks[i]

        results = run_batch_task_gen(task_gen, inputs, output_map=result2content_fn)
        output = "\n".join(
            f"{agent_names[i]}: {result}" for i, result in enumerate(results)
        )
        return AgentDoneTool(content=output)


def chat(model: str = "", sentence: str = None) -> None:

    cities_agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            name=CITY_AGENT_NAME,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or lm.OpenAIChatModel.GPT4o,
            ),
            system_message="""
            You'll receive a sentence. 
            Simply show the the list of cities in the sentence if any,
            as a comma-separated list, say nothing else.
            If no cities are found, say "NO CITIES".
            """,
        )
    )

    names_agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            name=NAME_AGENT_NAME,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or lm.OpenAIChatModel.GPT4o,
            ),
            system_message="""
            You'll receive a sentence. 
            Simply show the the list of names in the sentence if any,
            as a comma-separated list, say nothing else.
            If no names are found, say "NO NAMES".
            """,
        )
    )

    cities_task = lr.Task(cities_agent, interactive=False, single_round=True)
    names_task = lr.Task(names_agent, interactive=False, single_round=True)

    MyGlobalState.set_values(
        name_task_map={CITY_AGENT_NAME: cities_task, NAME_AGENT_NAME: names_task}
    )

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            name="MainAgent",
            llm=lm.OpenAIGPTConfig(
                chat_model=model or lm.OpenAIChatModel.GPT4o,
            ),
            system_message=f"""
            You'll receive a sentence. Your end-goal is to get the 
            list of cities and names mentioned in the sentence,
            BUT YOU DO NOT KNOW HOW TO EXTRACT THEM;
            you'll receive the help of {CITY_AGENT_NAME} and {NAME_AGENT_NAME} for this.
            You must use the TOOL `{MultiTaskTool.name()}` to send the sentence 
            to them.
            Once you receive the consolidated results,
            say "DONE" and show the list of cities and names.
            """,
        )
    )

    agent.enable_message(MultiTaskTool)

    task = lr.Task(agent, interactive=False, single_round=False)

    sentence = sentence or "Satoshi will meet Alice in New York and Bob in London"

    result = task.run(sentence)

    print(
        f"""
        [bold]Final Result:[/bold]
        {result}
        """
    )


if __name__ == "__main__":
    Fire(chat)
</file>

<file path="examples/basic/done_sequences_example.py">
#!/usr/bin/env python3
"""
Example demonstrating the new done_sequences feature in Langroid Tasks.

This feature allows you to specify sequences of events that trigger task completion,
providing more flexibility than simple done conditions.

You can use either:
1. DSL string patterns for convenience: "T, A" (tool then agent)
2. Full DoneSequence objects for more control

DSL Pattern Syntax:
- T = Any tool
- T[name] = Specific tool
- A = Agent response
- L = LLM response
- U = User response
- N = No response
- C[pattern] = Content matching regex

Note: Sequences use strict matching - events must occur consecutively in the message
chain without intervening messages. This ensures predictable behavior and efficient
matching.
"""

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import (
    AgentEvent,
    DoneSequence,
    EventType,
    Task,
    TaskConfig,
)
from langroid.agent.tool_message import ToolMessage


# Define a simple calculator tool
class CalculatorTool(ToolMessage):
    request: str = "calculator"
    purpose: str = "Perform arithmetic calculations"
    expression: str

    def handle(self) -> str:
        try:
            result = eval(self.expression)
            return f"The result is: {result}"
        except Exception as e:
            return f"Error: {str(e)}"


# Define a search tool
class SearchTool(ToolMessage):
    request: str = "search"
    purpose: str = "Search for information"
    query: str

    def handle(self) -> str:
        # Mock search implementation
        return f"Search results for '{self.query}': [Mock results here]"


def example0_dsl_syntax():
    """Example 0: Using DSL string patterns (recommended for simple cases)"""
    print("\n=== Example 0: DSL String Patterns ===")

    agent = ChatAgent(
        ChatAgentConfig(
            name="Assistant",
            system_message="""
            You are a helpful assistant with access to calculator and search tools.
            Use the appropriate tool when asked to calculate or search for something.
            """,
        )
    )
    agent.enable_message(CalculatorTool, use=True, handle=True)
    agent.enable_message(SearchTool, use=True, handle=True)

    # Using DSL string patterns - much more concise!
    config = TaskConfig(
        done_sequences=[
            "T, A",  # Any tool then agent response
            "T[calculator], A",  # Specific calculator tool
            "C[quit|exit|bye]",  # Content matching pattern
            "L, T, A, L",  # Complex sequence
        ]
    )

    _ = Task(agent, config=config)
    print("Task configured with multiple DSL patterns.")
    print(
        "Will complete on any of: tool use, calculator use, quit words, or L->T->A->L sequence"
    )
    # _ = task.run("What is 25 * 4?")
    # print(f"Final result: {result.content}")


def example1_tool_then_agent():
    """Example 1: Task completes after any tool is generated and handled by agent"""
    print("\n=== Example 1: Tool -> Agent Response ===")

    agent = ChatAgent(
        ChatAgentConfig(
            name="Assistant",
            system_message="""
            You are a helpful assistant with access to calculator and search tools.
            Use the appropriate tool when asked to calculate or search for something.
            """,
        )
    )
    agent.enable_message(CalculatorTool, use=True, handle=True)
    agent.enable_message(SearchTool, use=True, handle=True)

    # Task completes after: Tool -> Agent Response
    # Using DSL (recommended for simple patterns):
    # config = TaskConfig(done_sequences=["T, A"])

    # Using full syntax (for more control):
    config = TaskConfig(
        done_sequences=[
            DoneSequence(
                name="tool_handled",
                events=[
                    AgentEvent(event_type=EventType.TOOL),
                    AgentEvent(event_type=EventType.AGENT_RESPONSE),
                ],
            )
        ]
    )

    task = Task(agent, config=config)
    print("Task will complete after any tool is used and handled.")
    _ = task.run("What is 25 * 4?")
    # print(f"Final result: {_.content}")


def example2_specific_tool_sequence():
    """Example 2: Task completes only after specific tool (calculator) is used"""
    print("\n=== Example 2: Specific Tool Sequence ===")

    agent = ChatAgent(
        ChatAgentConfig(
            name="Assistant",
            system_message="""
            You help users with calculations and searches.
            Always use the appropriate tool.
            """,
        )
    )
    agent.enable_message(CalculatorTool, use=True, handle=True)
    agent.enable_message(SearchTool, use=True, handle=True)

    # Task completes only after calculator tool is used
    config = TaskConfig(
        done_sequences=[
            DoneSequence(
                name="calculation_done",
                events=[
                    AgentEvent(
                        event_type=EventType.SPECIFIC_TOOL, tool_name="calculator"
                    ),
                    AgentEvent(event_type=EventType.AGENT_RESPONSE),
                ],
            )
        ]
    )

    task = Task(agent, config=config)
    print("Task will complete only after calculator tool is used.")
    print("Try: 'Search for Python tutorials' (won't complete task)")
    print("Then try: 'Calculate 15 + 27' (will complete task)")
    _ = task.run()


def example3_conversation_pattern():
    """Example 3: Task completes after specific conversation pattern"""
    print("\n=== Example 3: Conversation Pattern ===")

    agent = ChatAgent(
        ChatAgentConfig(
            name="Assistant",
            system_message="""
            You are a step-by-step assistant. When asked to solve a problem:
            1. First acknowledge the request
            2. Then use the calculator tool
            3. Finally provide a summary of the result
            """,
        )
    )
    agent.enable_message(CalculatorTool, use=True, handle=True)

    # Task completes after: LLM -> Tool -> Agent -> LLM pattern
    config = TaskConfig(
        done_sequences=[
            DoneSequence(
                name="problem_solved",
                events=[
                    AgentEvent(event_type=EventType.LLM_RESPONSE),  # Acknowledgment
                    AgentEvent(event_type=EventType.TOOL),  # Calculator use
                    AgentEvent(event_type=EventType.AGENT_RESPONSE),  # Tool handled
                    AgentEvent(event_type=EventType.LLM_RESPONSE),  # Summary
                ],
            )
        ]
    )

    task = Task(agent, config=config)
    print("Task will complete after: acknowledgment -> tool use -> handling -> summary")
    _ = task.run(
        "I need to calculate the area of a rectangle with width 12 and height 8"
    )


def example4_multiple_completion_paths():
    """Example 4: Multiple ways to complete a task"""
    print("\n=== Example 4: Multiple Completion Paths ===")

    agent = ChatAgent(
        ChatAgentConfig(
            name="Assistant",
            system_message="""
            You help users with various tasks. 
            If they say 'quit' or 'exit', acknowledge and stop.
            Otherwise, help them with calculations or searches.
            """,
        )
    )
    agent.enable_message(CalculatorTool, use=True, handle=True)
    agent.enable_message(SearchTool, use=True, handle=True)

    # Multiple ways to complete the task
    config = TaskConfig(
        done_sequences=[
            # Path 1: User says quit/exit
            DoneSequence(
                name="user_quit",
                events=[
                    AgentEvent(
                        event_type=EventType.CONTENT_MATCH,
                        content_pattern=r"\b(quit|exit|bye|goodbye)\b",
                    ),
                ],
            ),
            # Path 2: Calculator tool used
            DoneSequence(
                name="calculation_done",
                events=[
                    AgentEvent(
                        event_type=EventType.SPECIFIC_TOOL, tool_name="calculator"
                    ),
                    AgentEvent(event_type=EventType.AGENT_RESPONSE),
                ],
            ),
            # Path 3: Two searches performed
            DoneSequence(
                name="double_search",
                events=[
                    AgentEvent(event_type=EventType.SPECIFIC_TOOL, tool_name="search"),
                    AgentEvent(event_type=EventType.AGENT_RESPONSE),
                    AgentEvent(event_type=EventType.SPECIFIC_TOOL, tool_name="search"),
                    AgentEvent(event_type=EventType.AGENT_RESPONSE),
                ],
            ),
        ]
    )

    task = Task(agent, config=config)
    print("Task can complete in 3 ways:")
    print("1. Say 'quit' or 'exit'")
    print("2. Use the calculator tool")
    print("3. Use the search tool twice")
    _ = task.run()


def example5_combining_with_existing_options():
    """Example 5: Combining done_sequences with done_if_tool"""
    print("\n=== Example 5: Combining with Existing Options ===")

    agent = ChatAgent(
        ChatAgentConfig(
            name="Assistant",
            system_message="You are a helpful assistant with tool access.",
        )
    )
    agent.enable_message(CalculatorTool, use=True, handle=True)

    # Combine done_sequences with done_if_tool
    config = TaskConfig(
        done_if_tool=True,  # Quick exit on any tool
        done_sequences=[
            # This won't be reached if done_if_tool triggers first
            DoneSequence(
                name="complex_pattern",
                events=[
                    AgentEvent(event_type=EventType.LLM_RESPONSE),
                    AgentEvent(event_type=EventType.LLM_RESPONSE),
                    AgentEvent(event_type=EventType.TOOL),
                ],
            )
        ],
    )

    task = Task(agent, config=config)
    print("Task will complete as soon as any tool is generated (done_if_tool=True)")
    _ = task.run("Calculate 5 + 5")


if __name__ == "__main__":
    print("Langroid Done Sequences Examples")
    print("=" * 50)

    # Run examples (comment out interactive ones if running all at once)
    example0_dsl_syntax()  # Show DSL syntax
    example1_tool_then_agent()
    # example2_specific_tool_sequence()  # Interactive
    # example3_conversation_pattern()    # May need specific LLM
    # example4_multiple_completion_paths()  # Interactive
    example5_combining_with_existing_options()

    print("\n" + "=" * 50)
    print("Examples completed!")
</file>

<file path="examples/basic/drug-outcomes.py">
"""
ADE (Adverse Drug Event) probability estimation task:

Given a pair of (Drug Category, Adverse Event), have the LLM generate an estimate
of the probability that the drug category is associated with an increased risk
of the adverse event.

Run this N times (without caching) to get statistics on the estimates.
Illustrates the use of `llm_response_batch`.

Default model is GPT4o, see how to specify alternative models below.

Example run:

python3 examples/basic/ drug-outcomes.py \
    --model litellm/claude-3-5-sonnet-20240620 --temp 0.1 \
    --pair "(Antibiotics, Acute Liver Injury)" --n 20 --reason true

Interesting models to try:
- gpt-4o (default)
- gpt-4
- litellm/claude-3-5-sonnet-20240620
- groq/llama3-70b-8192

See reference below for specific (DrugCategory, ADE) pairs to test.

References:
    - Guides to using Langroid with local and non-OpenAI models:
        https://langroid.github.io/langroid/tutorials/local-llm-setup/
        https://langroid.github.io/langroid/tutorials/non-openai-llms/
    - OMOP Ground Truth table of known Drug-ADE associations:
        (see page 16 for the table of Drug-ADE pairs)
        https://www.brookings.edu/wp-content/uploads/2012/04/OMOP-methods-review.pdf
"""

import re

import numpy as np
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.utils.configuration import settings

# Turn off cache retrieval, to get independent estimates on each run
settings.cache = False

MODEL = lm.OpenAIChatModel.GPT4o
TEMP = 0.1
PAIR = "(Antibiotics, Acute Liver Injury)"
N = 20
# should LLM include reasoning along with probability?
# (meant to test whether including reasoning along with the probability
# improves accuracy and/or variance of estimates)
REASON: bool = False


def extract_num(x: str) -> int:
    """
    Extracts an integer from a string that contains a number.

    Args:
        x (str): The input string containing the number.

    Returns:
        int: The extracted integer.

    Raises:
        ValueError: If no number is found in the expected format.
    """
    match = re.search(r"\d+", x)
    if match:
        return int(match.group(0))
    else:
        return -1


def main(
    model: str = MODEL,
    temp: float = TEMP,
    pair: str = PAIR,
    n: int = N,
    reason: bool = REASON,
):
    REASONING_PROMPT = (
        """
            IMPORTANT: Before showing your estimated probability, 
            you MUST show 2-3 sentences with your REASONING, and THEN give your 
            percent probability estimate in the range [0,100].
    """
        if reason
        else ""
    )

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=lm.OpenAIGPTConfig(
                temperature=temp,
                chat_model=model,
            ),
            name="ADE-Estimator",
            system_message=f"""
            You are a clinician with deep knowledge of Adverse Drug Events (ADEs) 
            of various drugs and categories of drugs.
            You will be given a (DRUG CATEGORY, ADVERSE OUTCOME) pair,
            you have to estimate the probability that this DRUG CATEGORY
            is associated with INCREASED RISK of the ADVERSE OUTCOME. 
    
            {REASONING_PROMPT}
                
            You must give your probability estimate as a SINGLE NUMBER e.g. 56, 
            which means 56%.             
            DO NOT GIVE A RANGE OF PROBABILITIES, ONLY A SINGLE NUMBER. 
            """,
        )
    )

    results = lr.llm_response_batch(
        agent,
        [pair] * n,
        # ["(Beta Blockers, Mortality after Myocardial Infarction)"]*20,
    )
    probs = [extract_num(r.content) for r in results]
    cached = [r.metadata.cached for r in results]
    n_cached = sum(cached)
    # eliminate negatives (due to errs)
    probs = [p for p in probs if p >= 0]
    mean = np.mean(probs)
    std = np.std(probs)
    std_err = std / np.sqrt(len(probs))
    hi = max(probs)
    lo = min(probs)
    print(f"Stats for {pair} with {model} temp {temp} reason {reason}:")
    print(
        f"N: {len(probs)} ({n_cached} cached ) Mean: {mean:.2f}, Std: {std:.2f}, StdErr:"
        f" {std_err:.2f}, min: {lo:.2f}, max: {hi:.2f}"
    )
    toks, cost = agent.llm.tot_tokens_cost()
    print(f"Tokens: {toks}, Cost: {cost:.2f}")


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/fn-call-local-numerical.py">
"""
Function-calling example using a local LLM, with ollama.

"Function-calling" refers to the ability of the LLM to generate
a structured response, typically a JSON object, instead of a plain text response,
which is then interpreted by your code to perform some action.
This is also referred to in various scenarios as "Tools", "Actions" or "Plugins".
See more here: https://langroid.github.io/langroid/quick-start/chat-agent-tool/

This script is designed to have a basic ChatAgent (powered by an Open-LLM)
engage in a multi-round conversation where the user may occasionally
ask for the "Polinsky transform" of a number, which requires the LLM to
use a `Polinsky` tool/function-call. This is a fictitious transform,
that simply does n => 3n + 1.
We intentionally use a fictitious transform rather than something like "square"
or "double" to prevent the LLM from trying to answer the question directly.

The challenging part here is getting the LLM to decide on an appropriate response
to a few different types of user messages:
- user asks a general question -> LLM should answer the question directly
- user asks for the Polinsky transform of a number -> LLM should use the Polinsky tool
- result from applying Polinsky transform -> LLM should present this to the user
- user (tool-handler) says there was a format error in using the Polinsky tool -> LLM
    should try this tool again

Many models quickly get confused in a multi-round conversation like this.
However (as of Sep 2024), `llama-3.1-70b` seems to do well here (we run this via groq).

Run like this --

python3 examples/basic/fn-call-local-numerical.py -m groq/llama-3.1-70b-versatile

or

python3 examples/basic/fn-call-local-numerical.py -m ollama/qwen2.5-coder:latest


(if the optional -m <model_name> is not provided, it defaults to GPT-4o).

See here for ways to set up a Local/Open LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import os
from typing import List, Optional

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_document import ChatDocument
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import ForwardTool
from langroid.language_models.openai_gpt import OpenAICallParams
from langroid.utils.configuration import settings

DEFAULT_LLM = lm.OpenAIChatModel.GPT4o


os.environ["TOKENIZERS_PARALLELISM"] = "false"

# (1) Define the desired fn-call as a ToolMessage via Pydantic.


class PolinskyTool(lr.agent.ToolMessage):
    """A fictitious number transformation tool. We intentionally use
    a fictitious tool rather than something like "square" or "double"
    to prevent the LLM from trying to answer the question directly.
    """

    request: str = "polinsky"
    purpose: str = (
        """
        To respond to user request for the Polinsky transform of a <number>.
        NOTE: ONLY USE THIS TOOL AFTER THE USER ASKS FOR A POLINSKY TRANSFORM. 
        """
    )
    number: int

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        # Used to provide few-shot examples in the system prompt
        return [
            cls(
                number=19,
            ),
            cls(
                number=5,
            ),
        ]


class MyChatAgent(lr.ChatAgent):
    def init_state(self) -> None:
        self.tool_expected = False

    def polinsky(self, msg: PolinskyTool) -> str:
        """Handle LLM's structured output if it matches Polinsky tool"""
        self.tool_expected = False
        result = msg.number * 3 + 1
        response = f"""
        SUCCESS! The Polinksy transform of {msg.number} is {result}.
        Present this result to the user, and ask what they need help with.
        """
        return response

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        self.tool_expected = True
        return super().llm_response(message)

    def user_response(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        self.tool_expected = False
        return super().user_response(msg)

    def handle_message_fallback(self, msg: str | ChatDocument) -> ForwardTool:
        """
        We end up here when there was no recognized tool msg from the LLM;
        In this case forward the message to the user using ForwardTool.
        """
        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            return ForwardTool(agent="User")


def app(
    m: str = DEFAULT_LLM,  # model name
    d: bool = False,  # debug
    nc: bool = False,  # no cache
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=16_000,  # for dolphin-mixtral
        max_output_tokens=100,
        params=OpenAICallParams(
            presence_penalty=0.8,
            frequency_penalty=0.8,
        ),
        temperature=0,
        stream=True,
        timeout=100,
    )

    # Recommended: First test if basic chat works with this llm setup as below:
    # Once this works, then you can try the rest of the example.
    #
    # agent = lr.ChatAgent(
    #     lr.ChatAgentConfig(
    #         llm=llm_cfg,
    #     )
    # )
    #
    # agent.llm_response("What is 3 + 4?")
    #
    # task = lr.Task(agent)
    # verify you can interact with this in a chat loop on cmd line:
    # task.run("Concisely answer some questions")

    # Define a ChatAgentConfig and ChatAgent
    config = lr.ChatAgentConfig(
        llm=llm_cfg,
        system_message="""
        You are an expert at deciding when to call 
        specified functions with the right syntax.
        You are very very CONCISE in your responses.
        
        Here is how you must respond to my messages:
        
        1. When I ask a general question, simply respond as you see fit.
            Example: 
                ME(User): "What is 3 + 4?"
                YOU(Assistant): "the answer is 7"
                
        2. When I ask to find the Polinksy transform of a number, 
            you  must use the `polinsky` function/tool
            to request the Polinsky transform of that number.
            Example:
                ME(User): "What is the Polinsky transform of 5?"
                YOU(Assistant): <polinsky tool request in JSON format>
                 
        3. When you receive a SUCCESS message with the result from the `polinsky` 
            tool, you must present the result to me in a nice way (CONCISELY), 
            and ask: 'What else can I help with?'
            Example:
                ME(User): "SUCCESS! The Polinksy transform of 5 is 16"
                YOU(Assistant): "The polinsky transform of 5 is 16. What else can I help with?"
                ME(User): "The answer is 16. What is the Polinsky transform of 19?"
                YOU(Assistant): <polinsky tool request in JSON format>
        4. If you receive an error msg when using the `polinsky` function/tool,
           you must try the function/tool again with the same number.
              Example:
               ME(User): "There was an error in your use of the polinsky tool:..."
               YOU(Assistant): <polinsky tool request in JSON format>
        """,
    )

    agent = MyChatAgent(config)

    # (4) Enable the Tool for this agent --> this auto-inserts JSON instructions
    # and few-shot examples into the system message
    agent.enable_message(PolinskyTool)

    # (5) Create task and run it to start an interactive loop
    task = lr.Task(agent, interactive=False)
    task.run("Can you help me with some questions?")


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/basic/fn-call-local-simple.py">
"""
Function-calling example using a local/remote open LLM.

"Function-calling" refers to the ability of the LLM to generate
a structured response, typically a JSON object, instead of a plain text response,
which is then interpreted by your code to perform some action.
This is also referred to in various scenarios as "Tools", "Actions" or "Plugins".
See more here: https://langroid.github.io/langroid/quick-start/chat-agent-tool/

Run like this (to run with llama-3.1-8b-instant via groq):

python3 examples/basic/fn-call-local-simple.py -m groq/llama-3.1-8b-instant

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


"""

import os
from typing import List

import fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_document import ChatDocument
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import FinalResultTool
from pydantic import BaseModel, Field
from langroid.utils.configuration import settings

# for best results:
DEFAULT_LLM = lm.OpenAIChatModel.GPT4o

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# (1) Define the desired structure via Pydantic.
# Here we define a nested structure for City information.
# The "Field" annotations are optional, and are included in the system message
# if provided, and help with generation accuracy.


class CityData(BaseModel):
    population: int = Field(..., description="population of city")
    country: str = Field(..., description="country of city")


class City(BaseModel):
    name: str = Field(..., description="name of city")
    details: CityData = Field(..., description="details of city")


# (2) Define the Tool class for the LLM to use, to produce the above structure.
class CityTool(lr.agent.ToolMessage):
    """Present information about a city"""

    request: str = "city_tool"
    purpose: str = """
    To present <city_info> AFTER user gives a city name,
    with all fields of the appropriate type filled out;
    """
    city_info: City = Field(..., description="information about a city")

    def handle(self) -> FinalResultTool:
        """Handle LLM's structured output if it matches City structure"""
        print("SUCCESS! Got Valid City Info")
        return FinalResultTool(answer=self.city_info)

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        # Used to provide few-shot examples in the system prompt
        return [
            cls(
                city_info=City(
                    name="San Francisco",
                    details=CityData(
                        population=800_000,
                        country="USA",
                    ),
                )
            )
        ]


def app(
    m: str = DEFAULT_LLM,  # model
    d: bool = False,  # pass -d to enable debug mode (see prompts etc)
    nc: bool = False,  # pass -nc to disable cache-retrieval (i.e. get fresh answers)
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=32000,  # set this based on model
        max_output_tokens=1000,
        temperature=0.2,
        stream=True,
        timeout=45,
    )

    # Recommended: First test if basic chat works with this llm setup as below:
    # Once this works, then you can try the rest of the example.
    #
    # agent = lr.ChatAgent(
    #     lr.ChatAgentConfig(
    #         llm=llm_cfg,
    #     )
    # )
    #
    # agent.llm_response("What is 3 + 4?")
    #
    # task = lr.Task(agent)
    # verify you can interact with this in a chat loop on cmd line:
    # task.run("Concisely answer some questions")

    # Define a ChatAgentConfig and ChatAgent

    config = lr.ChatAgentConfig(
        llm=llm_cfg,
        handle_llm_no_tool=f"""
            You FORGOT to use the TOOL/Function `{CityTool.name()}` 
            to present city info!
            """,
        system_message=f"""
        You will receive a city name, 
        and you must use the TOOL/FUNCTION `{CityTool.name()}` to generate/present
        information about the city. In other words, your response must 
        be a JSON string starting with `{{"request": "{CityTool.name()}", ...}}`
        """,
    )

    agent = lr.ChatAgent(config)

    # (4) Enable the Tool for this agent --> this auto-inserts JSON instructions
    # and few-shot examples (specified in the tool defn above) into the system message
    agent.enable_message(CityTool)

    # (5) Create task specialized to return City object
    task: City | None = lr.Task(agent, interactive=False)[City]

    while True:
        city = Prompt.ask("Enter a city name")
        if city in ["q", "x"]:
            break
        result: City | None = task.run(city)
        if result:
            print(f"City Info: {result}")
        else:
            print("No valid city info found.")


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/basic/intent-classifier.py">
"""
Agent-loop to classify the intent of a given text.

Run like this (--model is optional, defaults to GPT4o):

python3 examples/basic/intent-classifier.py --model groq/llama-3.1-8b-instant

Other ways to specify the model:
- gpt-4 (set OPENAI_API_KEY in your env or .env file)
- gpt-4o (ditto, set OPENAI_API_KEY)
- cerebras/llama3.1-70b (set CEREBRAS_API_KEY)

For more ways to use langroid with other LLMs, see:
- local/open LLMs: https://langroid.github.io/langroid/tutorials/local-llm-setup/
- non-OpenAPI LLMs: https://langroid.github.io/langroid/tutorials/non-openai-llms/
"""

from enum import Enum
from typing import List, Tuple

from fire import Fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ResultTool


class Intent(str, Enum):
    GREETING = "greeting"
    FAREWELL = "farewell"
    QUESTION = "question"
    STATEMENT = "statement"


class IntentTool(lr.ToolMessage):
    request: str = "intent_tool"
    purpose: str = """
        To classify the <intent> of a given text, into one of:
        - greeting
        - farewell
        - question
        - statement
        """

    intent: Intent

    @classmethod
    def examples(cls) -> List[lr.ToolMessage | Tuple[str, lr.ToolMessage]]:
        """Use these as few-shot tool examples"""
        return [
            cls(intent=Intent.GREETING),
            ("I want to classify this as a question", cls(intent=Intent.QUESTION)),
        ]

    def handle(self) -> ResultTool:
        """Handle the tool sent by LLM"""

        # ResultTool ends the task
        return ResultTool(intent=self.intent)

    def handle_message_fallback(self, message: lr.ChatDocument) -> str | None:
        """We end up here if the LLM did not send a Tool, so nudge it"""
        if (
            isinstance(message, lr.ChatDocument)
            and message.metadata.sender == lr.Entity.LLM
        ):
            return """
            You forgot to use the `intent_tool` to classify the intent.
            """


def main(model: str = ""):
    intent_tool_name = IntentTool.default_value("request")
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            name="Intent",
            llm=lm.OpenAIGPTConfig(chat_model=model or lm.OpenAIChatModel.GPT4o),
            use_functions_api=False,  # suppress OpenAI functions/tools
            use_tools=True,  # use langroid-native tools: works with ANY LLM
            system_message=f"""
            You are an astute INTENT CLASSIFIER: given any piece of text
            from the user, you are able to smartly infer their intent.
            Given such a piece of text, classify its intent into one of the following:
             - greeting
             - farewell
             - question
             - statement
            To present your classification, use the 
            `{intent_tool_name}` tool.
            
            ALWAYS use this tool to respond, do NOT say anything else.
            """,
        )
    )

    agent.enable_message(IntentTool)

    # create a task loop specialized to return an Intent
    task = lr.Task(agent=agent, interactive=False)[Intent]

    while True:
        text = Prompt.ask("Enter a text to classify its intent")
        intent = task.run(
            f"""
         Please classify the intent of this text, present your answer
         using the `{intent_tool_name}` tool:
         ----
         {text}
         ----
         """
        )

        print(f"Intent: {intent}")


if __name__ == "__main__":

    Fire(main)
</file>

<file path="examples/basic/multi-agent-medical.py">
"""
Credit to @burcusayin for contributing this example.

Run like this:

    python3 examples/basic/multi-agent-medical.py

or
    uv run examples/basic/multi-agent-medical.py

A two-agent system to answer medical questions that require a binary yes/no answer,
along with a `long_answer` explanation. The agents consist of:

- Chief Physician (CP) agent who is in charge of the final binary decision
    and explanation.
- Physician Assistant (PA) agent who is consulted by the CP; The CP may ask a
  series of questions to the PA, and once the CP decides they have sufficient
  information, they will return their final decision using a structured tool message.

The system is run over 445 medical questions from this dataset:
https://huggingface.co/datasets/burcusayin/pubmedqa_binary_with_plausible_gpt4_long_answers

In each row of this dataset, there is a QUESTION, and a final_decision
which we use as reference to compare the system-generated final decision.
"""

import logging

import datasets
import pandas as pd
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.task import TaskConfig
from langroid.agent.tools.orchestration import ForwardTool, ResultTool
from pydantic import BaseModel, Field
from langroid.utils.configuration import settings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
lr.utils.logging.setup_colored_logging()

# MODEL = lm.OpenAIChatModel.GPT4o
MODEL = "ollama/llama3:8b"

CP_NAME = "CP"
PA_NAME = "PA"


class ExpectedText(BaseModel):
    final_decision: str = Field(..., description="binary yes/no answer")
    long_answer: str = Field(..., description="explanation for the final decision")


class ExpectedTextTool(lr.ToolMessage):
    request: str = "expected_text_tool"
    purpose: str = """
    To write the final <expectedText> AFTER having a multi-turn discussion
    with the Assistant Agent, with all fields of the appropriate type filled out.
    """
    expectedText: ExpectedText

    def handle(self) -> ResultTool:
        """Handle LLM's structured output if it matches ExpectedText structure"""
        print("SUCCESS! Got Valid ExpectedText Info")

        return ResultTool(status="**DONE!**", expectedText=self.expectedText)

    @staticmethod
    def handle_message_fallback(
        agent: lr.ChatAgent, msg: str | lr.ChatDocument
    ) -> ForwardTool:
        """
        We end up here when there was no recognized tool msg from the LLM;
        In this case forward the message to the Assistant agent (PA) using ForwardTool.
        """
        if isinstance(msg, lr.ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            return ForwardTool(agent=PA_NAME)


# Define fixed system messages outside of the question-loop
# Pass each question as senior_task.run(question)

SENIOR_SYS_MSG = f"""You are Dr. X, the Chief Physician, collaborating with Dr. Y, your assistant.
                    Your task is to come up with concise answers to medical questions.
                    To make better decisions, when you receive a question, you should follow a TWO-PHASE procedure:

                    PHASE 1: Ask your assistant NATURAL LANGUAGE questions (NO TOOLS), which may span
                        MULTIPLE ROUNDS. ASK EXACTLY ONE QUESTION in each round. DO NOT ASK MULTIPLE QUESTIONS AT ONCE.
                        Avoid fabricating interactions or simulating dialogue with Dr. Y.
                        Instead, clearly articulate your questions or follow-ups, analyze Dr. Y's responses,
                        and use this information to guide your decision-making.
                    PHASE 2: Once you have gathered sufficient information, return your final decision
                        using the TOOL `{ExpectedTextTool.name()}`:
                        - `final_decision` should be your BINARY yes/no answer
                        - `long_answer` should provide a detailed explanation for your final decision.
                    DO NOT mention the TOOL to Dr. Y. It is your responsibility to write and submit the expectedText.
                    """

ASSISTANT_SYS_MSG = """You are Dr. Y, an assistant physician working under the supervision of Dr. X, the chief physician.
                            Your role is to respond to a medical question
                            by providing your initial evaluation, which will guide Dr. X
                            toward finalizing the answer. Dr X may ask you a series of questions, and you should respond
                            based on your expertise and the preceding discussion.
                            ### Instructions:
                            1. Ensure your evaluation is clear, precise, and structured to facilitate an informed discussion.
                            2. In each round of the discussion, limit yourself to a CONCISE message.
                        ### Process:
                        You will first receive a message from Dr. X, asking for your initial assessment.
                        Afterward, you can follow up in each discussion round to collaboratively refine the answer.
                        """


class ChatManager:
    def __init__(
        self,
        d: bool = False,  # pass -d to enable debug mode (see prompts etc)
        nc: bool = False,  # pass -nc to disable cache-retrieval (i.e. get fresh answer)
    ):
        settings.debug = d
        settings.cache = not nc

        self.ass_lm_config = lm.OpenAIGPTConfig(
            chat_model=MODEL,
            chat_context_length=1040_000,
            seed=42,
        )
        self.ass_agent = lr.ChatAgent(
            lr.ChatAgentConfig(
                name=PA_NAME,
                llm=self.ass_lm_config,
                system_message=ASSISTANT_SYS_MSG,
            ),
        )
        # no need for the DiscussionTextTool
        # self.ass_agent.enable_message(DiscussionTextTool)
        self.senior_lm_config = lm.OpenAIGPTConfig(
            chat_model=MODEL,
            chat_context_length=1040_000,
            seed=42,
        )
        self.senior_agent = lr.ChatAgent(
            lr.ChatAgentConfig(
                llm=self.senior_lm_config,
                name=CP_NAME,
                system_message=SENIOR_SYS_MSG,
            ),
        )
        self.senior_agent.enable_message(ExpectedTextTool)

    def start_chat(
        self, question: str
    ) -> ExpectedText:  # this is our main function to start the chat
        task_config = TaskConfig(inf_loop_cycle_len=0)
        self.ass_task = lr.Task(
            self.ass_agent,
            llm_delegate=True,
            interactive=False,
            single_round=True,
            config=task_config,
        )

        self.senior_task = lr.Task(
            self.senior_agent,
            llm_delegate=True,
            interactive=False,
            single_round=False,
            config=task_config,
        )[
            ResultTool
        ]  # specialize task to strictly return ResultTool or None

        self.senior_task.add_sub_task(self.ass_task)
        response_tool: ResultTool | None = self.senior_task.run(
            question, turns=100
        )  # dialogues usually take less than 70 turns

        if response_tool is None:
            print(
                """
                RETURNED ANSWER DOES NOT HAVE A TOOL! LLM DID NOT FORMAT THE DISCHARGE TEXT!!!
                """
            )
            return ExpectedText(final_decision="unknown", long_answer="null")
        else:
            print("ResultTool has been received successfully!!!")
            print(response_tool.expectedText)
            return response_tool.expectedText


if __name__ == "__main__":
    chatAgent = ChatManager()

    pubmed_ds = pd.DataFrame(
        datasets.load_dataset(
            "burcusayin/pubmedqa_binary_with_plausible_gpt4_long_answers"
        )["test"]
    )
    model_responses = []
    nrows = len(pubmed_ds)
    print(f"Processing {nrows} questions")
    for i, row in enumerate(pubmed_ds.itertuples()):
        question = row.QUESTION
        reference_decision = row.final_decision
        print(f"QUESTION: {question}")
        response: ExpectedText = chatAgent.start_chat(question=question)
        model_responses.append(response)
        print(
            f"Got response {i}: {response.final_decision}, reference: {reference_decision}"
        )
        cont = Prompt.ask("Continue? (y/n)", default="y")
        if cont.lower() != "y":
            break
</file>

<file path="examples/basic/multi-agent-return-result.py">
"""
3-agent system where Main task has subtasks that are able to directly return final
task result, "short-circuiting" the flow.

main_task has sub-tasks even_task and odd_task.

- main_task receives a number, simply passes it on.
- even_task can only handle even number N, returns N/2 as final result, 
    else passes it on.
- odd_task can only handle odd number N, returns 3N+1 as final result, 
    else passes it on.
"""

import langroid as lr
from langroid.agent.tools.orchestration import FinalResultTool

main_agent = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Main",
        system_message="Whatever number you receive, simply repeat it",
    )
)


class MyFinalResultTool(FinalResultTool):
    request: str = "my_final_result_tool"
    purpose: str = "To present the final result of the exercise"
    _allow_llm_use: bool = True

    answer: int  # could of course be str if answer is text


my_final_result_tool = MyFinalResultTool.default_value("request")

even_agent = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Even",
        system_message=f"""
        - If you receive an even number, return half of it using the 
          TOOL `{my_final_result_tool}` with `answer` set to your answer.
        - Otherwise simply repeat the number
        """,
    )
)

odd_agent = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Odd",
        system_message=f"""
        - If you receive an odd number N, return 3N+1 using the
          TOOL `{my_final_result_tool}` with `answer` set to your answer.
        - Otherwise simply repeat the number        
        """,
    )
)


even_agent.enable_message(MyFinalResultTool)
odd_agent.enable_message(MyFinalResultTool)

# set up main_task to return a result of type MyFinalResultTool
main_task = lr.Task(main_agent, interactive=False)[MyFinalResultTool]
even_task = lr.Task(even_agent, interactive=False)
odd_task = lr.Task(odd_agent, interactive=False)
main_task.add_sub_task([even_task, odd_task])

result = main_task.run("3")
assert isinstance(result, MyFinalResultTool)
assert result.answer == 10

result = main_task.run("4")
assert isinstance(result, MyFinalResultTool)
assert result.answer == 2

# we can also pass in an integer value

result = main_task.run(15)
assert isinstance(result, MyFinalResultTool)
assert result.answer == 46

result = main_task.run(16)
assert isinstance(result, MyFinalResultTool)
assert result.answer == 8
</file>

<file path="examples/basic/multi-agent-round-table.py">
"""
Toy example where 3 agents concurrently respond to the current message,
and the current message is updated to the response of one such responder.

Run like this:

python3 examples/basic/multi-agent-round-table.py

"""

import langroid as lr
from langroid.agent.batch import run_batch_task_gen
from langroid.utils.constants import NO_ANSWER

agent1 = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="agent1",
        system_message=f"""
        You are a simple number transformer, follow this rule:
        - If you see a number ending in 0,1, or 2, respond with a random 3-digit number.
        - Otherwise, respond saying: {NO_ANSWER}
        """,
    )
)
task1 = lr.Task(agent1, interactive=False, single_round=True)

agent2 = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="agent2",
        system_message=f"""
        You are a simple number transformer, follow this rule:
        - If you see a number ending in 3,4, or 5, respond with a random 3-digit number.
        - Otherwise, respond saying: {NO_ANSWER}
        """,
    )
)
task2 = lr.Task(agent2, interactive=False, single_round=True)


agent3 = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="agent3",
        system_message=f"""
        You are a simple number transformer, follow this rule:
        - If you see a number ending in 6,7,8 or 9, respond with a random 3-digit number.
        - Otherwise, respond saying: {NO_ANSWER}
        """,
    )
)
task3 = lr.Task(agent3, interactive=False, single_round=True)

tasks = [task1, task2, task3]


def task_gen(i):
    return tasks[i]


# kickoff with n = 412
n = 412
# run for 10 rounds
for _ in range(10):
    print("n = ", n)
    inputs = [n] * 3
    results = run_batch_task_gen(task_gen, inputs)
    # find which result is not NO_ANSWER
    for i, r in enumerate(results):
        if r.content != NO_ANSWER:
            n = int(r.content)
            print(f"agent{i+1} responded with {n}")
            break
</file>

<file path="examples/basic/multi-agent-triage.py">
"""
3-agent student assistant system:

- Triage agent: routes questions to the appropriate agent
- Course Agent: answers questions about courses
- Finance Agent: answers questions about finances

Illustrates use of AgentDoneTool, ForwardTool

Run like this (if --model is omitted, it defaults to the GPT-4o model):

python3 examples/basic/multi-agent-triage.py --model groq/llama-3.1-70b-versatile


"""

import os
from typing import Optional

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.tools.orchestration import (
    AgentDoneTool,
    ForwardTool,
    SendTool,
)
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig
from langroid.parsing.urls import find_urls
from langroid.vector_store.qdrantdb import QdrantDBConfig

os.environ["TOKENIZERS_PARALLELISM"] = "false"

forward_tool_name = ForwardTool.default_value("request")


class FinanceAnswerTool(lr.ToolMessage):
    request: str = "finance_answer_tool"
    purpose: str = "Present the <answer> to a question about finances"

    answer: str

    def handle(self) -> SendTool:
        return SendTool(to="User", content=self.answer)


class CoursesAnswerTool(lr.ToolMessage):
    request: str = "courses_answer_tool"
    purpose: str = "Present the <answer> to a question about courses"

    answer: str

    def handle(self) -> SendTool:
        return SendTool(to="User", content=self.answer)


def main(model: str = ""):
    class TriageAgent(lr.ChatAgent):
        def init_state(self) -> None:
            # self.expecting_course_answer = False
            # self.expecting_finance_answer = False
            super().init_state()
            self.llm_responded = False

        def user_response(
            self,
            msg: Optional[str | ChatDocument] = None,
        ) -> Optional[ChatDocument]:
            self.llm_responded = False
            return super().user_response(msg)

        def llm_response(
            self, message: Optional[str | ChatDocument] = None
        ) -> Optional[ChatDocument]:
            self.llm_responded = True
            return super().llm_response(message)

        def handle_message_fallback(
            self, msg: str | ChatDocument
        ) -> str | ChatDocument | lr.ToolMessage | None:
            """Handle any non-tool msg"""
            if self.llm_responded:
                self.llm_responded = False
                # LLM generated non-tool msg => send to user
                content = msg.content if isinstance(msg, ChatDocument) else msg
                return SendTool(to="User", content=content)

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        max_output_tokens=200,
        chat_context_length=16_000,
    )

    triage_agent = TriageAgent(
        lr.ChatAgentConfig(
            name="Triage",
            llm=llm_config,
            system_message=f"""
            You are a helpful assistant to students at a university. 
            
            Students may ask about the following TYPES of questions and you must handle 
            each TYPE as specified below:
            
            - (a) COURSES:
                - use the TOOL: `{forward_tool_name}` to forward the 
                    question to the "Courses" agent
            - (b) FINANCES (student loans, scholarships, tuition, dining plans, etc)
                - use the TOOL: `{forward_tool_name}` to forward the
                    question to the "Finance" agent
            - (c) OTHER questions not specific to the university:
                - attempt to answer these based on your own knowledge, 
                  otherwise admit you don't know.
            
            Start by greeting the user and asking them what they need help with.
            """,
        )
    )
    triage_agent.enable_message(ForwardTool)
    triage_agent.enable_message(
        [FinanceAnswerTool, CoursesAnswerTool],
        use=False,
        handle=True,
    )

    triage_task = lr.Task(triage_agent, interactive=False)

    parsing_config = ParsingConfig(  # modify as needed
        chunk_size=200,  # aim for this many tokens per chunk
        overlap=50,  # overlap between chunks
        max_chunks=10_000,
        # aim to have at least this many chars per chunk when
        # truncating due to punctuation
        min_chunk_chars=50,
        discard_chunk_chars=5,  # discard chunks with fewer than this many chars
        n_neighbor_ids=5,  # num chunk IDs to store on either side of each chunk
        pdf=PdfParsingConfig(
            # NOTE: PDF parsing is extremely challenging, and each library
            # has its own strengths and weaknesses.
            # Try one that works for your use case.
            # See here for available alternatives:
            # https://github.com/langroid/langroid/blob/main/langroid/parsing/parser.py
            library="pymupdf4llm",
        ),
    )

    class CoursesAgent(lr.agent.special.DocChatAgent):
        def llm_response(
            self,
            message: None | str | ChatDocument = None,
        ) -> Optional[ChatDocument]:
            answer = super().llm_response(message)
            if answer is None:
                return None
            return self.create_llm_response(
                tool_messages=[
                    AgentDoneTool(tools=[CoursesAnswerTool(answer=answer.content)])
                ]
            )

    course_url = "https://csd.cmu.edu/cs-and-related-undergraduate-courses"

    courses_agent = CoursesAgent(
        config=lr.agent.special.DocChatAgentConfig(
            name="Courses",
            llm=llm_config,
            doc_paths=[course_url],  # contents will be ingested into vecdb
            vecdb=QdrantDBConfig(
                collection_name="courses",
                replace_collection=True,
                storage_path=".qdrantdb/data/",
            ),
            parsing=parsing_config,
            n_neighbor_chunks=3,
            n_similar_chunks=5,
            n_relevant_chunks=5,
        )
    )

    courses_task = lr.Task(courses_agent, interactive=False, single_round=True)

    finance_url = "https://www.cmu.edu/sfs/tuition/index.html"
    all_finance_urls = find_urls(finance_url, max_links=20, max_depth=3)

    class FinanceAgent(lr.agent.special.DocChatAgent):
        def llm_response(
            self,
            message: None | str | ChatDocument = None,
        ) -> Optional[ChatDocument]:
            answer = super().llm_response(message)
            if answer is None:
                return None
            return self.create_llm_response(
                tool_messages=[
                    AgentDoneTool(tools=[FinanceAnswerTool(answer=answer.content)])
                ]
            )

    finance_agent = FinanceAgent(
        config=lr.agent.special.DocChatAgentConfig(
            name="Finance",
            llm=llm_config,
            doc_paths=all_finance_urls,  # contents will be ingested into vecdb
            vecdb=QdrantDBConfig(
                collection_name="finances",
                replace_collection=True,
                storage_path=".qdrantdb/data/",
            ),
            parsing=parsing_config,
            n_neighbor_chunks=3,
            n_similar_chunks=5,
            n_relevant_chunks=5,
        )
    )

    finance_task = lr.Task(finance_agent, interactive=False, single_round=True)

    triage_task.add_sub_task([courses_task, finance_task])

    triage_task.run()


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/oai-asst-chat.py">
"""
The most basic chatbot example, using an OpenAssistant agent,
powered by the OpenAI Assistant API.

Run like this:

python3 examples/basic/oai-asst-chat.py
"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

from langroid.agent.openai_assistant import OpenAIAssistant, OpenAIAssistantConfig
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


@app.command()
def chat() -> None:
    print(
        """
        [blue]Welcome to the basic chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    default_sys_msg = "You are a helpful assistant. Be concise in your answers."

    sys_msg = Prompt.ask(
        "[blue]Tell me who I am. Hit Enter for default, or type your own\n",
        default=default_sys_msg,
    )

    config = OpenAIAssistantConfig(
        system_message=sys_msg,
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
    )
    agent = OpenAIAssistant(config)
    task = Task(agent)

    task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/oai-code-chat.py">
"""
The most basic use of code-interpreter, using an OpenAssistant agent,
powered by the OpenAI Assistant API's code-interpreter tool.

Run like this:

python3 examples/basic/oai-code-chat.py
"""

import tempfile

import typer
from dotenv import load_dotenv
from rich import print
from rich.prompt import Prompt

from langroid.agent.openai_assistant import (
    AssistantTool,
    OpenAIAssistant,
    OpenAIAssistantConfig,
    ToolType,
)
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.parsing.url_loader import URLLoader
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


@app.command()
def chat() -> None:
    print(
        """
        [blue]Welcome to the basic chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    default_sys_msg = "You are a helpful assistant. Be concise in your answers."

    sys_msg = Prompt.ask(
        "[blue]Tell me who I am. Hit Enter for default, or type your own\n",
        default=default_sys_msg,
    )

    path = Prompt.ask("Enter a URL or file path, or hit enter if no files")
    if path:
        # if path is a url, use UrlLoader to get text as a document
        if path.startswith("http"):
            text = URLLoader([path]).load()[0].content
            # save text to a temp file
            with tempfile.NamedTemporaryFile(
                mode="w", suffix=".txt", delete=False
            ) as f:
                f.write(text)
                f.close()
                # get the filename
                path = f.name

    config = OpenAIAssistantConfig(
        system_message=sys_msg,
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
    )
    agent = OpenAIAssistant(config)
    agent.add_assistant_tools([AssistantTool(type=ToolType.CODE_INTERPRETER)])
    if path:
        agent.add_assistant_files([path])

    task = Task(agent)

    task.run(
        """
        Help me with some questions, 
        using the CODE INTERPRETER tool, and any uploaded files as needed.
        """
    )


if __name__ == "__main__":
    app()
</file>

<file path="examples/basic/plan-subtasks.py">
"""
Planner agent receives a math calculation expression from user,
involving + - * / ops, with possible parentheses. Planner has no math abilities,
so it needs to create a plan of elementary operations to compute the result,
and send each step to the appropriate helper agent, who will return the result.

Run like this:

python3 examples/basic/plan-subtasks.py

When it waits for user input, try asking things like:

- (10 + 2)/6 - 1
- 3*(4+1) - 3

"""

import langroid as lr
from langroid.utils.constants import AT, DONE, NO_ANSWER

planner = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Planner",
        system_message=f"""
        User will give you a math calculation, but you have no math abilities.
        However you are a great planner, so your task is to do two things:
        
        1. CREATE a PLAN of what 
          sequence of ELEMENTARY operations (ONLY add/subtract, multiply/divide) need
          to performed, in order to compute what the user asked for.
        2. EMIT the needed operations, ONE BY ONE, and wait for the answer from
            each, before emitting the next operation. Since you cannot directly
            calculate these, you will have to SEND the needed operations to 
            specific helpers, as follows:
            
            * Send Multiplication operation to `Multiplier`
            * Send Add operation to `Adder`
            * Send Subtract operation to `Subtractor`
            * Send Divide operation to `Divider`
            
            To clarify who you are sending the message to, preface your message with
            {AT}<helper_name>, e.g. "{AT}Multiplier multiply with 5" 
            
            When you have the final answer, say {DONE} and show it.
            
            At the START, ask the user what they need help with, 
            address them as "{AT}user"
            
        EXAMPLE: 
        ============
        User: please calculate (4*5 + 1)/3
        Assistant (You): 
            PLAN: 
                1. multiply 4 with 5
                2. add 1 to the result
                3. divide result by 3
            {AT}Multiplier multiply 4 with 5
            [... wait for result, then show your NEW PLAN and send a new request]
            and so on.                         
                        
        """,
    )
)

adder = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Adder",
        system_message=f"""
        If you receive an Add request, return the result,
        otherwise say {NO_ANSWER}.
        """,
    )
)

multiplier = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Multiplier",
        system_message=f"""
        If you receive a Multiply request, return the result,
        otherwise say {NO_ANSWER}.
        """,
    )
)

subtractor = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Subtractor",
        system_message=f"""
        If you receive a Subtraction request, return the result,
        otherwise say {NO_ANSWER}.
        """,
    )
)

divider = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="Divider",
        system_message=f"""
        If you receive a Division request, return the result,
        otherwise say {NO_ANSWER}.
        """,
    )
)


task_config = lr.TaskConfig(addressing_prefix=AT)
planner_task = lr.Task(planner, interactive=False, config=task_config)
adder_task = lr.Task(adder, interactive=False, single_round=True)
multiplier_task = lr.Task(multiplier, interactive=False, single_round=True)
divider_task = lr.Task(divider, interactive=False, single_round=True)
subtractor_task = lr.Task(subtractor, interactive=False, single_round=True)

planner_task.add_sub_task([adder_task, multiplier_task, divider_task, subtractor_task])


planner_task.run()
</file>

<file path="examples/basic/planner-workflow-simple.py">
"""
Illustrates a Planner agent orchestrating a multi-step workflow by using tools that
invoke other specialized agents.

- The PlannerAgent is instructed to first increment a number by 3, and then
  multiply the result by 8.
- To do this, it repeatedly uses two tools: `IncrementTool` and `DoublingTool`.
- The key idea is that these tools are stateful: their `handle_async` methods
  don't perform the simple math themselves, but instead run other `Task` objects
  (`increment_task`, `doubling_task`).
- These tasks are handled by simple, specialized agents (`IncrementAgent`,
  `DoublingAgent`) that only know how to perform a single, small step.

This example showcases a powerful pattern where a high-level agent delegates complex
sub-processes to other agents via the tool mechanism.

Run like this from the repo root, once you are in a virtual environment with
langroid installed:

    uv run examples/basic/planner-workflow-simple.py

To use a different model, for example, run like this:

    uv run examples/basic/planner-workflow-simple.py --model gpt-4.1-mini

"""

import logging

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import DoneTool
from pydantic import Field

logger = logging.getLogger(__name__)
MODEL = lm.OpenAIChatModel.GPT4_1


class IncrementAgentConfig(lr.ChatAgentConfig):
    name: str = "Incrementer"
    system_message: str = "Given a number, return the next number"


class DoublingAgentConfig(lr.ChatAgentConfig):
    name: str = "Doubler"
    system_message: str = "Given a number, return the number multiplied by 2"


async def main(model: str = ""):

    increment_agent = lr.ChatAgent(
        IncrementAgentConfig(
            llm=lm.OpenAIGPTConfig(
                chat_model=model or MODEL,
                async_stream_quiet=False,
            )
        )
    )
    increment_task = lr.Task(
        increment_agent,
        interactive=False,
        single_round=True,
    )

    doubling_agent = lr.ChatAgent(
        DoublingAgentConfig(
            llm=lm.OpenAIGPTConfig(
                chat_model=model or MODEL,
                async_stream_quiet=False,
            )
        )
    )

    doubling_task = lr.Task(
        doubling_agent,
        interactive=False,
        single_round=True,
    )

    class IncrementTool(lr.ToolMessage):
        request: str = "increment_tool"
        purpose: str = "To increment a <number> by 1"
        number: int = Field(..., description="The number (int) to Increment")

        async def handle_async(self) -> str:
            # stateful tool: handler runs the increment_task
            result = await increment_task.run_async(f"{self.number}")
            return result.content

    class DoublingTool(lr.ToolMessage):
        request: str = "doubling_tool"
        purpose: str = "To double a <number>"
        number: int = Field(..., description="The number (int) to Double")

        async def handle_async(self) -> str:
            # stateful tool: handler runs the doubling_task
            result = await doubling_task.run_async(self.number)
            return result.content

    class PlannerConfig(lr.ChatAgentConfig):
        name: str = "Planner"
        handle_llm_no_tool: str = "You FORGOT to use one of your TOOLs!"
        llm: lm.OpenAIGPTConfig = lm.OpenAIGPTConfig(
            chat_model=model or MODEL,
            async_stream_quiet=False,
        )
        system_message: str = f"""
        You are a Planner in charge of PROCESSING the user's input number
        (an integer) through a SEQUENCE of two steps:
        
        1. Increment the number by 3 -- use the `{IncrementTool.name()}` tool,
            as many times as needed, until the number is incremented by 3.
        2. Multiply the number by 8 -- use the `{DoublingTool.name()}` tool,
            as many times as needed, until the number is multiplied by 8.
            
        Note That even though these tasks sound trivial, you cannot and must not do them 
        yourself. You must use the tools as many times as needed for each step and then 
        proceed to the next step. 
        
        CRITICAL: You must call ONE TOOL only and wait for its result, 
        and then call another tool. 
        NEVER EVER call multiple tools at the same time.  
        
        Once you are done, use the TOOL `{DoneTool.name()}` to return the final result.
        """

    planner = lr.ChatAgent(PlannerConfig())

    planner.enable_message([IncrementTool, DoublingTool, DoneTool])

    planner_task = lr.Task(planner, interactive=False)

    result = await planner_task.run_async("Process this number: 3")
    assert "48" in result.content, f"Expected 48, got {result.content}"


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/planner-workflow-spawn.py">
"""
Illustrates a Planner agent orchestrating a multi-step workflow by using the `TaskTool`
to dynamically spawn specialized sub-agents for each step.

- The PlannerAgent is instructed to first increment a number by 3, and then
  multiply the result by 8.
- To do this, it uses the `TaskTool` to dynamically create and run sub-tasks.
- For the incrementing part, it spawns a simple `IncrementAgent` three times.
- For the multiplication part, it spawns a simple `DoublingAgent` three times.

This example showcases a powerful pattern where a high-level agent can delegate
complex sub-processes to dynamically created, specialized agents without needing
them to be pre-defined in the main script.

Run like this from the repo root:

    uv run examples/basic/planner-workflow-spawn.py

To use a different model, for example gpt-4-turbo, run:

    uv run examples/basic/planner-workflow-spawn.py --model gpt-4-turbo

"""

import logging

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import DoneTool, ResultTool
from langroid.agent.tools.task_tool import TaskTool

logger = logging.getLogger(__name__)
MODEL = lm.OpenAIChatModel.GPT4_1


async def main(model: str = ""):

    class PlannerConfig(lr.ChatAgentConfig):
        name: str = "Planner"
        handle_llm_no_tool: str = "You FORGOT to use one of your TOOLs!"
        llm: lm.OpenAIGPTConfig = lm.OpenAIGPTConfig(
            chat_model=model or MODEL,
        )
        system_message: str = f""" 
        You are a Planner that has ZERO knowledge about MATH/ARITHMETIC!
        
        Your job is to process a number given by the user through a sequence of 2 steps:

        1.  **Increment the number by 3.**
        2.  **Multiply the resulting number by 8.**

        HOWEVER, you CANNOT do these steps yourself, so you instead 
        MUST use the `{TaskTool.name()}` to spawn a sub-agent for one of
        the following tasks as you see fit:
        
        - Increment a given number by 1
        - Double a given number
        
        The sub-agent can use "gpt-4.1-mini" as the model,
        and does not need any tools enabled.
        
        Keep track of the intermediate results.

        Once you have the final result, you MUST use the `{DoneTool.name()}` to return it.
        """

    planner = lr.ChatAgent(PlannerConfig())

    planner.enable_message([TaskTool, DoneTool])

    planner_task = lr.Task(planner, interactive=False)

    # Initial number is 3.
    # After incrementing 3 times: 3 + 3 = 6
    # After doubling 3 times: 6 * 2 * 2 * 2 = 48
    result = await planner_task.run_async("Process this number: 3")
    assert "48" in result.content, f"Expected 48, got {result.content}"


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/planner-workflow.py">
"""
Task: Process a number through a sequence of two steps:
- Burify: increment the number by 3
- Tonify: multiply the number by 4

Planner Agent oversees the process, using two worker agents:
- BurifyAgent: handles the Burify step
- TonifyAgent: handles the Tonify step

Planner checks intermediate results and provides feedback to worker agents,
until their step is complete, before proceeding to the next step.

Run like this from repo root (omit `-m` to use default model gpt-4.1-mini):

    uv run examples/basic/planner-workflow.py -m gpt-4.1-mini
"""

import logging
from typing import List

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool
from pydantic import Field

logger = logging.getLogger(__name__)
MODEL = lm.OpenAIChatModel.GPT4_1_MINI


class BurifyTool(lr.ToolMessage):
    request: str = "burify_tool"
    purpose: str = "To apply the 'Burify' process to a <number>"
    number: int = Field(..., description="The number (int) to Burify")

    def handle(self) -> str:
        # stateless tool: handler used in BurifyAgent
        return f"Burify this number: {self.number}"


class TonifyTool(lr.ToolMessage):
    request: str = "tonify_tool"
    purpose: str = "To apply the 'Tonify' process to a <number>"
    number: int = Field(..., description="The number (int) to Tonify")

    def handle(self) -> str:
        # stateless tool: handler used in TonifyAgent
        return f"Tonify this number: {self.number}"


class BurifyCheckTool(lr.ToolMessage):
    request: str = "burify_check_tool"
    purpose: str = "To check if the Burify process is complete"
    number: int = Field(..., description="The number (int) to check")
    original_number: int = Field(
        ...,
        description="The original number (int) given to the BurifyAgent",
    )

    def handle(self) -> str:
        # stateless tool
        if self.number == self.original_number + 3:
            return AcceptTool(result=self.number)
        else:
            return BurifyRevisionTool(
                feedback="Burify is NOT complete! Please try again.",
                recipient="Burify",
            )


class TonifyCheckTool(lr.ToolMessage):
    request: str = "tonify_check_tool"
    purpose: str = "To check if the Tonify process is complete"
    number: int = Field(..., description="The number (int) to check")
    original_number: int = Field(
        ...,
        description="The original number (int) given to the TonifyAgent",
    )

    def handle(self):
        # stateless tool
        if self.number == self.original_number * 4:
            return AcceptTool(result=self.number)
        else:
            return TonifyRevisionTool(
                feedback="Tonify is NOT complete! Please try again.",
                recipient="Tonify",
            )


class BurifyRevisionTool(lr.ToolMessage):
    request: str = "burify_revision_tool"
    purpose: str = "To give <feedback> to the  'BurifyAgent' on their Burify Attempt"
    feedback: str = Field(..., description="Feedback for the BurifyAgent")

    def handle(self):
        return f"""
        Below is feedback on your attempt to Burify: 
        <Feedback>
        {self.feedback}
        </Feedback>
        Please try again!
        """


class TonifyRevisionTool(lr.ToolMessage):
    request: str = "tonify_revision_tool"
    purpose: str = "To give <feedback> to the  'TonifyAgent' on their Tonify Attempt"
    feedback: str = Field(..., description="Feedback for the TonifyAgent")

    def handle(self):
        return f"""
        Below is feedback on your attempt to Tonify: 
        <Feedback>
        {self.feedback}
        </Feedback>
        Please try again!
        """


class BurifySubmitTool(lr.ToolMessage):
    request: str = "burify_submit_tool"
    purpose: str = "To submit the result of an attempt of the Burify process"
    result: int = Field(..., description="The result (int) to submit")

    def handle(self):
        return AgentDoneTool(content=str(self.result))


class TonifySubmitTool(lr.ToolMessage):
    request: str = "tonify_submit_tool"
    purpose: str = "To submit the result of an attempt of the Tonify process"
    result: int = Field(..., description="The result (int) to submit")

    def handle(self):
        return AgentDoneTool(content=str(self.result))


class AcceptTool(lr.ToolMessage):
    request: str = "accept_tool"
    purpose: str = "To accept the result of the 'Burify' or 'Tonify' process"
    result: int


class PlannerConfig(lr.ChatAgentConfig):
    name: str = "Planner"
    steps: List[str] = ["Burify", "Tonify"]
    handle_llm_no_tool: str = "You FORGOT to use one of your TOOLs!"
    system_message: str = f"""
    You are a Planner in charge of PROCESSING a given integer through
    a SEQUENCE of 2 processing STEPS, which you CANNOT do by yourself, but you must
    rely on WORKER AGENTS who will do these for you:
    - Burify - will be done by the BurifyAgent
    - Tonify - will be done by the TonifyAgent
    
    In order to INITIATE each process, you MUST use the appropriate TOOLs:
    - `{BurifyTool.name()}` to Burify the number (the tool will be handled by the BurifyAgent)
    - `{TonifyTool.name()}` to Tonify the number (the tool will be handled by the TonifyAgent)
    
    Each of the WORKER AGENTS works like this:
    - The Agent will ATTEMPT a processing step, using the number you give it.
    - You will VERIFY whether the processing step is COMPLETE or NOT
         using the CORRESPONDING CHECK TOOL:
         - check if the Burify step is complete using the `{BurifyCheckTool.name()}`
         - check if the Tonify step is complete using the `{TonifyCheckTool.name()}`
    - If the step is NOT complete, you will ask the Agent to try again,
        by using the CORRESPONDING Revision TOOL where you can include your FEEDBACK: 
        - `{BurifyRevisionTool.name()}` to revise the Burify step
        - `{TonifyRevisionTool.name()}` to revise the Tonify step
    - If you determine (see below) that the step is COMPLETE, you MUST
        use the `{AcceptTool.name()}` to ACCEPT the result of the step.    
    """


class PlannerAgent(lr.ChatAgent):
    current_step: int
    current_num: int
    original_num: int

    def __init__(self, config: PlannerConfig):
        super().__init__(config)
        self.config: PlannerConfig = config
        self.current_step = 0
        self.current_num = 0

    def burify_tool(self, msg: BurifyTool) -> str:
        """Handler of BurifyTool: uses/updates Agent state"""
        self.original_num = msg.number
        logger.warning(f"Planner handled BurifyTool: {self.current_num}")

        return ForwardTool(agent="Burify")

    def tonify_tool(self, msg: TonifyTool) -> str:
        """Handler of TonifyTool: uses/updates Agent state"""
        self.original_num = msg.number
        logger.warning(f"Planner handled TonifyTool: {self.current_num}")

        return ForwardTool(agent="Tonify")

    def accept_tool(self, msg: AcceptTool) -> str:
        """Handler of AcceptTool: uses/updates Agent state"""
        curr_step_name = self.config.steps[self.current_step]
        n_steps = len(self.config.steps)
        self.current_num = msg.result
        if self.current_step == n_steps - 1:
            # last step -> done
            return AgentDoneTool(content=str(self.current_num))

        self.current_step += 1
        next_step_name = self.config.steps[self.current_step]
        return f"""
            You have ACCEPTED the result of the {curr_step_name} step.
            Your next step is to apply the {next_step_name} process
            to the result of the {curr_step_name} step, which is {self.current_num}.
            So use a TOOL to initiate the {next_step_name} process!
            """


class BurifyAgentConfig(lr.ChatAgentConfig):
    name: str = "Burify"
    handle_llm_no_tool: str = f"You FORGOT to use the TOOL `{BurifySubmitTool.name()}`!"
    system_message: str = f"""
    You will receive an integer from your supervisor, to apply
    a process Burify to it, which you are not quite sure how to do,
    but you only know that it involves INCREMENTING the number by 1 a few times
    (but you don't know how many times).
    When you first receive a number to Burify, simply return the number + 1.
    If this is NOT sufficient, you will be asked to try again, and 
    you must CONTINUE to return your last number, INCREMENTED by 1.
    To send your result, you MUST use the TOOL `{BurifySubmitTool.name()}`. 
    """


class TonifyAgentConfig(lr.ChatAgentConfig):
    name: str = "Tonify"
    handle_llm_no_tool: str = f"You FORGOT to use the TOOL `{TonifySubmitTool.name()}`!"
    system_message: str = f"""
    You will receive an integer from your supervisor, to apply
    a process Tonify to it, which you are not quite sure how to do,
    but you only know that it involves MULTIPLYING the number by 2 a few times
    (and you don't know how many times).
    When you first receive a number to Tonify, simply return the number * 2.
    If this is NOT sufficient, you will be asked to try again, and 
    you must CONTINUE to return your last number, MULTIPLIED by 2.
    To send your result, you MUST use the TOOL `{TonifySubmitTool.name()}`.
    """


def main(model: str = ""):
    planner = PlannerAgent(
        PlannerConfig(
            llm=lm.OpenAIGPTConfig(
                chat_model=model or MODEL,
            )
        ),
    )

    planner.enable_message(
        [
            BurifyRevisionTool,
            TonifyRevisionTool,
        ],
        use=True,  # LLM allowed to generate
        handle=False,  # agent cannot handle
    )

    planner.enable_message(  # can use and handle
        [
            AcceptTool,
            BurifyCheckTool,
            TonifyCheckTool,
            BurifyTool,
            TonifyTool,
        ]
    )

    burifier = lr.ChatAgent(
        BurifyAgentConfig(
            llm=lm.OpenAIGPTConfig(
                chat_model=model or MODEL,
            )
        )
    )
    burifier.enable_message(
        [
            BurifyTool,
            BurifyRevisionTool,
        ],
        use=False,  # LLM cannot generate
        handle=True,  # agent can handle
    )
    burifier.enable_message(BurifySubmitTool)

    tonifier = lr.ChatAgent(
        TonifyAgentConfig(
            llm=lm.OpenAIGPTConfig(
                chat_model=model or MODEL,
            )
        )
    )

    tonifier.enable_message(
        [
            TonifyTool,
            TonifyRevisionTool,
        ],
        use=False,  # LLM cannot generate
        handle=True,  # agent can handle
    )
    tonifier.enable_message(TonifySubmitTool)

    planner_task = lr.Task(planner, interactive=False)
    burifier_task = lr.Task(burifier, interactive=False)
    tonifier_task = lr.Task(tonifier, interactive=False)

    planner_task.add_sub_task(
        [
            burifier_task,
            tonifier_task,
        ]
    )

    # Buify(5) = 5+3 = 8; Tonify(8) = 8*4 = 32
    result = planner_task.run("Sequentially all processes to this number: 5")
    assert "32" in result.content, f"Expected 32, got {result.content}"


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/python-code-exec-tool.py">
"""
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 contextlib
import io

from fire import Fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ResultTool
from pydantic import Field


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)
</file>

<file path="examples/basic/schedule-extract.py">
"""
Extract schedule/availability information from unstructured text.

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
"""

from typing import Dict, List, Literal, Tuple

from fire import Fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import FinalResultTool
from pydantic import BaseModel, Field


class Slot(BaseModel):
    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")


class DaySchedule(BaseModel):
    """
    A class to represent a day's schedule.
    """

    slots: List[Slot] = Field(..., description="List of time slots for the day")


Weekday = Literal["Mon", "Tue", "Wed", "Thu", "Fri"]


class Availability(BaseModel):
    """
    A class to represent schedule information.
    """

    week_availability: Dict[Weekday, DaySchedule] = Field(
        ...,
        description="""
        Dictionary mapping weekday to DaySchedule,
        where weekday is one of "Mon", "Tue", "Wed", "Thu", "Fri"
        """,
    )


class AvailabilityTool(lr.ToolMessage):
    request: str = "availability_tool"
    purpose: str = """
        To present the available slots from a piece of text.
    """
    availabilities: Availability

    @classmethod
    def examples(cls) -> List["lr.ToolMessage" | Tuple[str, "lr.ToolMessage"]]:
        """
        Example of how to use the tool.
        """
        return [
            (
                """
                I figured out that the availability is 10am-4pm on Mon and Wed at 
                home, and 3-4pm on Monday at daycare
                """,
                cls(
                    availabilities=Availability(
                        week_availability={
                            "Mon": DaySchedule(
                                slots=[
                                    Slot(
                                        start_time="10:00",
                                        end_time="16:00",
                                        location="home",
                                    ),
                                    Slot(
                                        start_time="15:00",
                                        end_time="16:00",
                                        location="daycare",
                                    ),
                                ]
                            ),
                            "Wed": DaySchedule(
                                slots=[
                                    Slot(
                                        start_time="10:00",
                                        end_time="16:00",
                                        location="home",
                                    )
                                ]
                            ),
                        }
                    )
                ),
            )
        ]

    def handle(self) -> str:
        """
        This method is called when the tool is invoked.
        It processes the input and returns the availability information.
        """
        # Here, we would implement the logic to extract availability information
        # from the input text. For this example, we'll just return a placeholder.
        print("Successfully extracted availability information.")
        print(self.availabilities.model_dump_json(indent=2))
        return FinalResultTool(avails=self.availabilities)


def make_schedule_task(model: str = ""):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o_MINI,
    )
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_config,
            system_message=f"""
            You are an expert at figuring out schedules from unstructured text.
            You will be given a string that represents availability information.
            Your task is to figure out the available slots and present this info
            using the TOOL `{AvailabilityTool.name()}`, with the `week_availability` 
            field set to a dictionary showing the available slots for certain days
            of the week if any. The string you will get may contain MULTIPLE 
            availabilities for the same day, but at different locations. 
            You have to present the availability information in the `availabilities`
            field, as an Availability object, which is a dictionary mapping
            the day of the week to a DaySchedule object, which is a list of
            Slot objects. The Slot object contains the start time of the slot,
            the duration of the slot in minutes, and the location of the slot.
            """,
        )
    )
    agent.enable_message(AvailabilityTool)
    task = lr.Task(agent, interactive=False, restart=True)[Availability]
    return task


def main(model: str = ""):
    task = make_schedule_task(model)
    while True:
        sched = Prompt.ask("Enter your schedule text")
        avails = task.run(sched, allow_restart=True)
        print(avails)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/text-to-structured.py">
"""
Function-calling example using a local LLM, with ollama.

"Function-calling" refers to the ability of the LLM to generate
a structured response, typically a JSON object, instead of a plain text response,
which is then interpreted by your code to perform some action.
This is also referred to in various scenarios as "Tools", "Actions" or "Plugins".
See more here: https://langroid.github.io/langroid/quick-start/chat-agent-tool/

Run like this (to run with llama-3.1-8b-instant via groq):

python3 examples/basic/text-to-structured.py -m groq/llama-3.1-8b-instant

Other models to try it with:
- ollama/qwen2.5-coder
- ollama/qwen2.5


See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/


"""

import json
import os
from typing import List, Literal

import fire
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import ResultTool
from pydantic import BaseModel, Field
from langroid.utils.configuration import settings

# for best results:
DEFAULT_LLM = lm.OpenAIChatModel.GPT4o

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# (1) Define the desired structure via Pydantic.
# The "Field" annotations are optional, and are included in the system message
# if provided, and help with generation accuracy.


class Wifi(BaseModel):
    name: str


class HomeSettings(BaseModel):
    App: List[str] = Field(..., description="List of apps found in text")
    wifi: List[Wifi] = Field(..., description="List of wifi networks found in text")
    brightness: Literal["low", "medium", "high"] = Field(
        ..., description="Brightness level found in text"
    )


# (2) Define the Tool class for the LLM to use, to produce the above structure.
class HomeAutomationTool(lr.agent.ToolMessage):
    """Tool to extract Home Automation structure from text"""

    request: str = "home_automation_tool"
    purpose: str = """
    To extract <home_settings> structure from a given text.
    """
    home_settings: HomeSettings = Field(
        ..., description="Home Automation settings from given text"
    )

    def handle(self) -> str:
        """Handle LLM's structured output if it matches HomeAutomationTool structure"""
        print(
            f"""
            SUCCESS! Got Valid Home Automation Settings:
            {json.dumps(self.home_settings.model_dump(), indent=2)}
            """
        )
        return ResultTool(settings=self.home_settings)

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        # Used to provide few-shot examples in the system prompt
        return [
            (
                """
                    I have extracted apps Spotify and Netflix, 
                    wifi HomeWifi, and brightness medium
                    """,
                cls(
                    home_settings=HomeSettings(
                        App=["Spotify", "Netflix"],
                        wifi=[Wifi(name="HomeWifi")],
                        brightness="medium",
                    )
                ),
            )
        ]


def app(
    m: str = DEFAULT_LLM,  # model
    d: bool = False,  # pass -d to enable debug mode (see prompts etc)
    nc: bool = False,  # pass -nc to disable cache-retrieval (i.e. get fresh answers)
):
    settings.debug = d
    settings.cache = not nc
    # create LLM config
    llm_cfg = lm.OpenAIGPTConfig(
        chat_model=m or DEFAULT_LLM,
        chat_context_length=4096,  # set this based on model
        max_output_tokens=100,
        temperature=0.2,
        stream=True,
        timeout=45,
    )

    tool_name = HomeAutomationTool.default_value("request")
    config = lr.ChatAgentConfig(
        llm=llm_cfg,
        system_message=f"""
        You are an expert in extracting home automation settings from text.
        When user gives a piece of text, use the TOOL `{tool_name}`
        to present the extracted structured information.
        """,
    )

    agent = lr.ChatAgent(config)

    # (4) Enable the Tool for this agent --> this auto-inserts JSON instructions
    # and few-shot examples (specified in the tool defn above) into the system message
    agent.enable_message(HomeAutomationTool)

    # (5) Create task and run it to start an interactive loop
    # Specialize the task to return a ResultTool object
    task = lr.Task(agent, interactive=False)[ResultTool]

    # set up a loop to extract Home Automation settings from text
    while True:
        text = Prompt.ask("[blue]Enter text (or q/x to exit)")
        if not text or text.lower() in ["x", "q"]:
            break
        result = task.run(text)
        assert isinstance(result, ResultTool)
        assert isinstance(result.settings, HomeSettings)


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/basic/tool-custom-handler.py">
"""
Short example of using `_handler` attribute in ToolMessage to define
custom name for `Agent` tool handler.

Run like this:

python3 examples/basic/tool-custom-handler.py

"""

import requests

import langroid as lr
from pydantic import Field


class CountryLanguageTool(lr.agent.ToolMessage):
    request: str = "country_language_tool"
    purpose: str = "To determine <language> spoken in specific country."
    country_name: str = Field(..., description="country name")
    _handler: str = "country_tools_handler"


class CountryPopulationTool(lr.agent.ToolMessage):
    request: str = "country_population_tool"
    purpose: str = "To determine <population> of specific country."
    country_name: str = Field(..., description="country name")
    _handler: str = "country_tools_handler"


class CountryAreaTool(lr.agent.ToolMessage):
    request: str = "country_area_tool"
    purpose: str = "To determine <area> of specific country."
    country_name: str = Field(..., description="country name")
    _handler: str = "country_tools_handler"


class AssistantAgent(lr.ChatAgent):
    def country_tools_handler(self, tool: lr.agent.ToolMessage):
        response = requests.get(
            f"https://restcountries.com/v3.1/name/{tool.country_name}", timeout=5
        )
        if not response.ok:
            return "invalid country name"

        try:
            data = response.model_dump_json()[0]
        except (ValueError, IndexError):
            return "invalid response"

        match tool.request:
            case "country_language_tool":
                language = ", ".join(data["languages"].values())
                return language
            case "country_population_tool":
                population_millions = data["population"] / 1e6
                return f"{population_millions:.1f} million people"
            case "country_area_tool":
                area_sq_km = data["area"] / 1e6
                return f"{area_sq_km:.1f} million sq. km"

        return "invalid tool name"


def make_assistant_task() -> lr.Task:
    llm_config = lr.language_models.OpenAIGPTConfig(
        temperature=0.2, max_output_tokens=250
    )

    assistant_config = lr.ChatAgentConfig(
        system_message="""
        You are a helpful assistant helping users with country-related questions.

        You know answers to the following questions:
          - what is the <language> spoken in specific country?
          - what is <population> of specific country?
          - what is <areay> of specific country?

        Ask user for the country name and information that he is interested in.
        Then use the appropriate tool to find the answer.
        """,
        llm=llm_config,
    )

    assistant_agent = AssistantAgent(assistant_config)
    assistant_agent.enable_message(CountryLanguageTool)
    assistant_agent.enable_message(CountryPopulationTool)
    assistant_agent.enable_message(CountryAreaTool)

    assistant_task = lr.Task(agent=assistant_agent, interactive=True)
    return assistant_task


if __name__ == "__main__":
    task = make_assistant_task()
    task.run()
</file>

<file path="examples/basic/tool-extract-short-example.py">
"""
Short example of using Langroid ToolMessage to extract structured info from a passage,
and perform computation on it.

Run like this (omit --model to default to GPT4o):

python3 examples/basic/tool-extract-short-example.py --model deepseek/deepseek-reasoner

or

uv run examples/basic/tool-extract-short-example.py --model deepseek/deepseek-reasoner

"""

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ResultTool
from pydantic import BaseModel, Field


# desired output structure
class CompanyInfo(BaseModel):
    name: str = Field(..., description="name of company")
    shares: int = Field(..., description="shares outstanding of company")
    price: float = Field(..., description="price per share of company")


# tool definition based on this
class CompanyInfoTool(lr.agent.ToolMessage):
    request: str = "company_info_tool"  # agent method that handles this tool
    purpose: str = (
        "To extract <company_info> from a passage and compute market-capitalization."
    )
    company_info: CompanyInfo

    @classmethod
    def examples(cls):
        """Examples that will be compiled to few-shot examples for the LLM.
        Illustrating two types of examples below:
        - example instance
        - (thought, example) tuple
        """
        return [
            # Example 1: just the instance
            cls(company_info=CompanyInfo(name="IBM", shares=1.24e9, price=140.15)),
            # Example 2: (thought, instance) tuple
            (
                "I want to extract and present company info from the passage",
                cls(
                    company_info=CompanyInfo(name="Apple", shares=16.82e9, price=149.15)
                ),
            ),
        ]

    def handle(self) -> ResultTool:
        """Handle LLM's structured output if it matches CompanyInfo structure.
        This suffices for a "stateless" tool.
        If the tool handling requires agent state, then
        instead of this `handle` method, define a `company_info_tool`
        method in the agent.
        Since this method is returning a  ResultTool,
        the task of this agent will be terminated,
        with this tool T appearing in the result ChatDocument's `tool_messages` list.
        """
        mkt_cap = self.company_info.shares * self.company_info.price
        return ResultTool(
            market_cap=mkt_cap,
            info=self.company_info,
            comment="success",  # arbitrary undeclared fields allowed
        )


# define agent, attach the tool


def main(model: str = ""):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_config,
            system_message=f"""
            Use the TOOL `{CompanyInfoTool.name()}` 
            tool to extract company information from a passage
            and compute market-capitalization.
            """,
        )
    )

    agent.enable_message(CompanyInfoTool)

    # define and run task on a passage about some company

    task = lr.Task(agent, interactive=False)

    print(
        """
        [blue]Welcome to the company info extractor!
        Write a sentence containing company name, shares outstanding and share price,
        and the Agent will use a tool/function extract the info in structured form,
        and the tool-handler will compute the market-cap.[/blue]
        """
    )

    while True:
        statement = Prompt.ask(
            """
            Enter a sentence containing company name, 
            shares outstanding and share price, or 
            hit enter to use default sentence.
            """,
            default="""
            Qualcomm has shares outstanding of 1.12 billion and a 
            price per share of $217.09.
            """,
        )
        result = task.run(statement)
        if result is None:
            print("Tool-call failed, try again.")
            continue
        # note the result.tool_messages will be a list containing
        # an obj of type FinalResultTool, so we can extract fields from it.
        company_result = result.tool_messages[0]
        assert isinstance(company_result, ResultTool)
        assert isinstance(company_result.info, CompanyInfo)

        info = company_result.info
        mktcap = company_result.market_cap
        assert company_result.comment == "success"
        print(
            f"""
            Found company info: {info} and market cap: {mktcap}
            """
        )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/basic/xml_tool.py">
"""
Example of defining a variant of an existing tool, but inheriting from XMLToolMessage,
to have the LLM use XML rather than JSON to generate the tool.

This will not work with built-in functions/tools of OpenAI,
so in the `ChatAgentConfig` , you have to set the following to ensure
that Langroid's built-in XML Tool calls are activated:
- `use_functions_api = False`
- `use_tools = True`

Run like this (--model is optional, defaults to GPT4o):

python3 examples/basic/xml_tool.py --model groq/llama-3.1-8b-instant
"""

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import SendTool
from langroid.agent.xml_tool_message import XMLToolMessage
from pydantic import Field


class XMLSendTool(SendTool, XMLToolMessage):
    """
    Variant of SendTool, using XML rather than JSON.
    """

    request: str = "xml_send_tool"
    purpose: str = """
        To send <content> to an entity/agent identified in the <to> field.
        """

    content: str = Field(
        ...,
        description="The content to send",
        verbatim=True,  # enforces content enclosed within CDATA block in xml.
    )
    to: str


xml_send_tool_name = XMLSendTool.default_value("request")


def main(model: str = ""):
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    alice = lr.ChatAgent(
        lr.ChatAgentConfig(
            name="Alice",
            llm=llm_config,
            use_functions_api=False,
            use_tools=True,
            system_message=f"""
            Whatever number you receive, send it to Bob using the  
            `{xml_send_tool_name}` tool.
            """,
        )
    )

    bob = lr.ChatAgent(
        lr.ChatAgentConfig(
            name="Bob",
            llm=llm_config,
            use_functions_api=False,
            use_tools=True,
            system_message=f"""
            Whatever number you receive, add 1 to it and send 
            the result back to Alice
            using the `{xml_send_tool_name}` tool.
            """,
        )
    )

    alice.enable_message(XMLSendTool)
    bob.enable_message(XMLSendTool)

    # specialize alice_task to return an int
    alice_task = lr.Task(alice, interactive=False)[int]
    bob_task = lr.Task(bob, interactive=False)

    alice_task.add_sub_task(bob_task)

    result = alice_task.run("5", turns=6)
    assert result == 7


if __name__ == "__main__":
    fire.Fire(main)
</file>

<file path="examples/chainlit/non-callback/chat-doc-qa-no-callback.py">
"""
Basic single-agent chat example, without streaming.

DEPCRECATED: Script kept only for reference. Best way is to use ChainlitAgentCallbacks,
as in chat-doc-qa.py

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-doc-qa-no-callback.py

Note, to run this with a local LLM, you can click the settings symbol
on the left of the chat window and enter the model name, e.g.:

ollama/mistral:7b-instruct-v0.2-q8_0

or

local/localhost:8000/v1"

depending on how you have set up your local LLM.

For more on how to set up a local LLM to work with Langroid, see:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import chainlit as cl

import langroid.language_models as lm
import langroid.parsing.parser as lp
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig


async def setup_agent() -> None:
    model = cl.user_session.get("settings", {}).get("ModelName")
    print(f"Using model: {model}")
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=16_000,  # adjust based on model
        timeout=90,
    )

    config = DocChatAgentConfig(
        n_query_rephrases=0,
        hypothetical_answer=False,
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=0,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        llm=llm_config,
        parsing=lp.ParsingConfig(  # modify as needed
            splitter=lp.Splitter.TOKENS,
            chunk_size=300,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=lp.PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )
    agent = DocChatAgent(config)
    cl.user_session.set("agent", agent)
    file = cl.user_session.get("file")
    msg = cl.Message(content=f"Processing `{file.name}`...", disable_feedback=True)
    await msg.send()
    agent.ingest_doc_paths([file.path])
    msg.content = f"Processing `{file.name}` done. Ask questions!"
    await msg.update()


@cl.on_settings_update
async def update_agent(settings):
    cl.user_session.set("settings", settings)
    await setup_agent()


@cl.on_chat_start
async def on_chat_start():
    await cl.ChatSettings(
        [
            cl.input_widget.TextInput(
                id="ModelName",
                label="Model Name (Default GPT4-Turbo)",
                default="",
            )
        ]
    ).send()

    # get file
    files = None
    # Wait for the user to upload a file
    while files is None:
        files = await cl.AskFileMessage(
            content="Please upload a text file to begin!",
            accept=["text/plain"],
            max_size_mb=20,
            timeout=180,
        ).send()

    file = files[0]
    print(f"got file: {file.name}")
    cl.user_session.set("file", file)
    await setup_agent()


@cl.on_message
async def on_message(message: cl.Message):
    agent: DocChatAgent = cl.user_session.get("agent")
    msg = cl.Message(content="")

    # need to do this since DocChatAgent does not have an async version of llm_response
    response = await cl.make_async(agent.llm_response)(message.content)
    msg.content = response.content
    await msg.send()
</file>

<file path="examples/chainlit/non-callback/chat-no-callback.py">
"""
Basic single-agent chat example, without streaming.

DEPCRECATED: Script kept only for reference.
The better way is shown in chat-agent.py or chat-task.py, which uses callbacks.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-no-callback.py
"""

import chainlit as cl

import langroid as lr


@cl.on_chat_start
async def on_chat_start():
    sys_msg = "You are a helpful assistant. Be concise in your answers."
    config = lr.ChatAgentConfig(
        system_message=sys_msg,
    )
    agent = lr.ChatAgent(config)
    cl.user_session.set("agent", agent)


@cl.on_message
async def on_message(message: cl.Message):
    agent: lr.ChatAgent = cl.user_session.get("agent")
    response = await agent.llm_response_async(message.content)
    msg = cl.Message(content=response.content)
    await msg.send()
</file>

<file path="examples/chainlit/non-callback/chat-search-no-callback.py">
"""
Basic single-agent chat example, using a web Search Tool, without streaming.

DEPCRECATED: Script kept only for reference. The better way is shown in
chat-search.py, which uses ChainlitTaskCallbacks.

- User asks a question
- LLM either responds directly or generates a Metaphor web search Tool/function-call
    - if Tool used:
         - Agent handler recognizes this tool and returns search results
         - LLM sees search results and composes a response.
- user asks another question


After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-search-no-callback.py
"""

import chainlit as cl

import langroid as lr
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool


@cl.step(name="LLM Response")
async def llm_response(msg: str) -> lr.ChatDocument:
    agent: lr.ChatAgent = cl.user_session.get("agent")
    response = await agent.llm_response_async(msg)
    return response


@cl.step(name="Agent Tool Handler")
async def agent_response(msg: lr.ChatDocument) -> lr.ChatDocument:
    agent: lr.ChatAgent = cl.user_session.get("agent")
    response = await agent.agent_response_async(msg)
    return response


@cl.on_chat_start
async def on_chat_start():
    tool_name = MetaphorSearchTool.default_value("request")
    sys_msg = f"""
        You are an astute, self-aware AI assistant, and you are adept at 
        responding to a user's question in one of two ways:
        - If you KNOW the answer from your own knowledge, respond directly.
        - OTHERWISE, request up to 5 results from a web search using 
          the `{tool_name}` tool/function-call.
          In this case you will receive the web search results, and you can 
          then compose a response to the user's question. 
    """
    config = lr.ChatAgentConfig(
        system_message=sys_msg,
    )
    agent = lr.ChatAgent(config)
    agent.enable_message(MetaphorSearchTool)
    cl.user_session.set("agent", agent)


@cl.on_message
async def on_message(message: cl.Message):
    agent: lr.ChatAgent = cl.user_session.get("agent")
    msg = cl.Message(content="")
    # expecting a tool here
    response = await llm_response(message.content)
    if agent.has_tool_message_attempt(response):
        search_results = await agent_response(response)
        response = await llm_response(search_results)
    msg.content = response.content
    await msg.send()
</file>

<file path="examples/chainlit/non-callback/chat-stream.py">
"""
DEPRECATED, not guaranteed to work: We are keeping this example for reference,
but do not use this as way to chat with streaming.
See chat-callback.py for the best way to do this
(i.e. use ChainlitAgentCallbacks when interacting directly an Agent,
or use ChainlitTaskCallbacks when interacting with a Task).

Basic single-agent chat example, with streaming,
using an older method, rather than the best way,
which is via callbacks, as in chat-callback.py.


After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-stream.py
"""

import asyncio
import re
import sys

import chainlit as cl

from langroid import ChatAgent, ChatAgentConfig
from langroid.utils.configuration import settings

settings.stream = True  # works if False as well


class ContinuousCaptureStream:
    """
    Capture stdout in a stream.
    This allows capturing of streaming output that would normally be printed to stdout,
    e.g. streaming tokens coming from OpenAI's API.
    """

    def __init__(self):
        self.content = ""
        self.new_content_event = asyncio.Event()
        self.is_finished = False  # Flag to indicate completion

    def write(self, data):
        self.content += data
        self.new_content_event.set()

    def flush(self):
        pass

    async def get_new_content(self):
        await self.new_content_event.wait()
        self.new_content_event.clear()
        new_content, self.content = self.content, ""
        return new_content

    def set_finished(self):
        self.is_finished = True
        self.new_content_event.set()  # T


def strip_ansi_codes(text):
    ansi_escape = re.compile(
        r"(?:\x1B[@-_]|[\x80-\x9A\x9C-\x9F]|[\x1A-\x1C\x1E-\x1F])+\[[0-?]*[ -/]*[@-~]"
    )
    return ansi_escape.sub("", text)


@cl.on_chat_start
async def on_chat_start():
    sys_msg = "You are a helpful assistant. Be concise in your answers."
    config = ChatAgentConfig(
        system_message=sys_msg,
        show_stats=False,
    )
    agent = ChatAgent(config)
    cl.user_session.set("agent", agent)


@cl.on_message
async def on_message(message: cl.Message):
    agent: ChatAgent = cl.user_session.get("agent")
    msg = cl.Message(content="")
    await msg.send()

    capture_stream = ContinuousCaptureStream()
    original_stdout = sys.stdout
    sys.stdout = capture_stream

    # Run response() in a separate thread or as a non-blocking call
    asyncio.create_task(run_response(agent, message, capture_stream))

    while not capture_stream.is_finished:
        new_output = await capture_stream.get_new_content()
        new_output = strip_ansi_codes(new_output)
        if new_output:
            await msg.stream_token(new_output)

    # Restore original stdout when done
    sys.stdout = original_stdout

    await msg.update()


async def run_response(agent: ChatAgent, message: cl.Message, stream):
    await agent.llm_response_async(message.content)
    stream.set_finished()
</file>

<file path="examples/chainlit/non-callback/chat-tool-no-callback.py">
"""
Basic single-agent chat example, using a Tool, without streaming.
DEPCRECATED: Script kept only for reference.
The better way is shown in chat-task-tool, which uses ChainlitTaskCallbacks.

- User enters a country
- LLM responds with a tool/function-call showing {country=country, capital=...}
- Agent handler recognizes this tool and returns plain text version of the tool result.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-tool-no-callback.py
"""

import chainlit as cl

import langroid as lr


class CapitalTool(lr.ToolMessage):
    request: str = "capital"
    purpose: str = "To present the capital of given <country>."
    country: str
    capital: str

    def handle(self) -> str:
        return f"""
        Success! LLM responded with a tool/function-call, with result:
        Capital of {self.country} is {self.capital}.
        """


@cl.step
async def llm_tool_call(msg: str) -> lr.ChatDocument:
    agent: lr.ChatAgent = cl.user_session.get("agent")
    response = await agent.llm_response_async(msg)
    return response


@cl.on_chat_start
async def on_chat_start():
    sys_msg = """
        You are an expert in country capitals.
        When user gives a country name, you should respond 
        with the capital of that country, using the `capital` tool/function-call.
    """
    config = lr.ChatAgentConfig(
        system_message=sys_msg,
    )
    agent = lr.ChatAgent(config)
    agent.enable_message(CapitalTool)
    cl.user_session.set("agent", agent)


@cl.on_message
async def on_message(message: cl.Message):
    agent: lr.ChatAgent = cl.user_session.get("agent")
    msg = cl.Message(content="")
    # expecting a tool here
    tool = await llm_tool_call(message.content)
    tool_result = await agent.agent_response_async(tool)
    msg.content = tool_result.content
    await msg.send()
</file>

<file path="examples/chainlit/non-callback/README.md">
## Chainlit examples without using Callbacks

These are all deprecated, but retaining here for reference.
The much better way to use Langroid Agents/Tasks with Chainlit is to use 
the `ChainlitAgentCallbacks` and `ChainlitTaskCallbacks` classes.
</file>

<file path="examples/chainlit/books.txt">
Book Title: Crime and Redemption by Filidor Dostoyevski, released in 1877, offers a
riveting exploration of guilt, morality, and the possibility of spiritual rebirth.
Set against the bleak backdrop of 19th century Russia, it follows the tormented journey
of Rodion Romanovich Raskolnikov, a young man driven to murder and subsequently
haunted by his actions. Through Raskolnikov's story, Dostoyevski delves deep into the
human psyche, presenting a timeless narrative of human imperfection and the
redemptive power.

Book Title: The Siblings Karamazoff by Fyodar Dostoyevskiy, published in 1881,
weaves a complex narrative around the ethical battles and spiritual dilemmas
faced by the Karamazoff family. Set in the heart of Russia, it explores themes of faith,
doubt, and the nature of free will through the intersecting lives of three brothers,
each embodying different facets of humanity. Dostoyevskiy masterfully crafts a tale of
familial bonds, existential questioning, and the search for truth in a morally ambiguous
world.
</file>

<file path="examples/chainlit/chainlit.md">
# Welcome to Langroid 👋

![Langroid](public/langroid-card.png)

---
When it is your turn to enter a message, you can do one of two things:
- write `c` to tell the agent to continue,
    - This is provided as a safeguard against infinite loops, or to prevent a large 
    amount of text to be sent to the LLM (which can be costly + slow). 
    If you simply want to continue with normal operation, just enter c.
- write a response, question or feedback to the agent, depending on context.
</file>

<file path="examples/chainlit/chat-doc-qa.py">
"""
Document question-answering using RAG on a single file, using ChainlitAgentCallbacks.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-doc-qa.py

Note, to run this with a local LLM, you can click the settings symbol
on the left of the chat window and enter the model name, e.g.:

ollama/mistral:7b-instruct-v0.2-q8_0

or

local/localhost:8000/v1"

depending on how you have set up your local LLM.

For more on how to set up a local LLM to work with Langroid, see:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

from textwrap import dedent

import chainlit as cl

import langroid as lr
import langroid.parsing.parser as lp
from langroid.agent.callbacks.chainlit import (
    SYSTEM,
    add_instructions,
    get_text_files,
    make_llm_settings_widgets,
    setup_llm,
    update_llm,
)
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.utils.constants import NO_ANSWER


async def initialize_agent() -> None:
    await setup_llm()
    llm_config = cl.user_session.get("llm_config")
    config = DocChatAgentConfig(
        name="DocAgent",
        n_query_rephrases=0,
        hypothetical_answer=False,
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=0,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        llm=llm_config,
        parsing=lp.ParsingConfig(  # modify as needed
            splitter=lp.Splitter.TOKENS,
            chunk_size=300,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=lp.PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )
    agent = DocChatAgent(config)
    cl.user_session.set("agent", agent)


@cl.on_settings_update
async def on_update(settings):
    await update_llm(settings)
    await initialize_agent()


@cl.on_chat_start
async def on_chat_start():
    await add_instructions(
        title="Basic Doc-Question-Answering using RAG (Retrieval Augmented Generation).",
        content=dedent(
            """
        **Upload** a document (click the attachment button in the chat dialog) and ask questions.
        **Change LLM settings** by clicking the settings symbol on the left of the chat window.
        
        You can keep uploading more documents, and questions will be answered based on all documents.
        """
        ),
    )

    await make_llm_settings_widgets()

    cl.user_session.set("callbacks_inserted", False)
    await initialize_agent()


@cl.on_message
async def on_message(message: cl.Message):
    agent: DocChatAgent = cl.user_session.get("agent")
    file2path = await get_text_files(message)
    agent.callbacks.show_start_response(entity="llm")
    if len(file2path) > 0:
        n_files = len(file2path)
        waiting = cl.Message(
            author=SYSTEM, content=f"Received {n_files} files. Ingesting..."
        )
        await waiting.send()
        agent.ingest_doc_paths(list(file2path.values()))
        file_or_files = "file" if n_files == 1 else "files"
        file_list = "\n".join([f"- `{file}`" for file in file2path.keys()])
        waiting.content = dedent(
            f"""
            Ingested `{n_files}` {file_or_files}:
            {file_list}
            """
        )
        await waiting.update()

    if not cl.user_session.get("callbacks_inserted", False):
        # first time user entered a msg, so inject callbacks and display first msg
        lr.ChainlitAgentCallbacks(agent)

    # Note DocChatAgent has no llm_response_async,
    # so we use llm_response with make_async
    response: lr.ChatDocument | None = await cl.make_async(agent.llm_response)(
        message.content
    )
    if response.content.strip() == NO_ANSWER:
        # in this case there were no relevant extracts
        # and we never called the LLM, so response was not shown in UI,
        # hence we need to send it here
        # TODO: It is possible the LLM might have already responded with NO_ANSWER,
        # so we may be duplicating the response here.
        agent.callbacks.show_llm_response(content=NO_ANSWER)
</file>

<file path="examples/chainlit/chat-search-assistant-local.py">
"""
Chainlit version of examples/basic/chat-search-assistant-local.py,
with a minor change to enable Chainlit callbacks.
Tested and works ok nous-hermes2-mixtral, but may still have issues.
See that script for details.

You can specify a local model in a few different ways, e.g. `groq/llama3-70b-8192`
or `ollama/mistral` etc. See here how to use Langroid with local LLMs:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

Since chainlit does not take cmd line args in the normal way, you have to specify
the model via an environment variable, e.g. `MODEL=ollama/mistral` before the
script is run, e.g.

MODEL=ollama/mistral chainlit run  examples/chainlit/chat-search-assistant-local.py

Note - this is just an example of using an open/local LLM;
 it does not mean that this will work with ANY local LLM.

You may get good results using `groq/llama3-70b-8192` (see the above-linked guide
to using open/local LLMs with Langroid for more details).

"""

import os
from textwrap import dedent
from typing import List, Optional, Type

import chainlit as cl
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.callbacks.chainlit import add_instructions
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool
from langroid.utils.configuration import Settings, set_global


class QuestionTool(lr.ToolMessage):
    request: str = "question_tool"
    purpose: str = "Ask a SINGLE <question> that can be answered from a web search."
    question: str

    @classmethod
    def examples(cls) -> List[lr.ToolMessage]:
        return [
            cls(question="Which superconductor material was discovered in 2023?"),
            cls(question="What AI innovation did Meta achieve in 2024?"),
        ]


class FinalAnswerTool(lr.ToolMessage):
    request: str = "final_answer_tool"
    purpose: str = """
        Present the intermediate <steps> and 
        final <answer> to the user's original query.
        """
    steps: str
    answer: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            cls(
                steps="1. Man is mortal. 2. Plato was a man.",
                answer="Plato was mortal.",
            ),
            cls(
                steps="1. The moon landing was in 1969. 2. Kennedy was president "
                "during 1969.",
                answer="Kennedy was president during the moon landing.",
            ),
        ]


class FeedbackTool(lr.ToolMessage):
    request: str = "feedback_tool"
    purpose: str = "Provide <feedback> on the user's answer."
    feedback: str

    @classmethod
    def examples(cls) -> List["lr.ToolMessage"]:
        return [
            cls(feedback=""),
            cls(
                feedback="""
                The answer is invalid because the conclusion does not follow from the
                steps. Please check your reasoning and try again.
                """
            ),
        ]


class AssistantAgent(lr.ChatAgent):
    n_questions: int = 0  # how many questions in THIS round
    has_asked: bool = False  # has ANY question been asked
    original_query: str | None = None

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.USER:
            # either first query from user, or returned result from Searcher
            self.n_questions = 0  # reset search count

        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            if self.has_asked:
                return f"""
                You may have intended to use a tool, but your JSON format may be wrong.
                
                REMINDER: You must do one of the following:
                - If you are ready with the final answer to the user's ORIGINAL QUERY
                    [ Remember it was: {self.original_query} ],
                  then present your reasoning steps and final answer using the 
                  `final_answer_tool` in the specified JSON format.
                - If you still need to ask a question, then use the `question_tool`
                  to ask a SINGLE question that can be answered from a web search.
                """
            elif self.original_query is not None:
                return f"""
                You must ask a question using the `question_tool` in the specified format,
                to break down the user's original query: {self.original_query} into 
                smaller questions that can be answered from a web search.
                """

    def question_tool(self, msg: QuestionTool) -> str:
        self.n_questions += 1
        self.has_asked = True
        if self.n_questions > 1:
            # there was already a search, so ignore this one
            return ""
        # valid question tool: re-create it so Searcher gets it
        return msg.to_json()

    def final_answer_tool(self, msg: FinalAnswerTool) -> str:
        if not self.has_asked or self.n_questions > 1:
            # not yet asked any questions, or LLM is currently asking
            # a question (and this is the second one in this turn, and so should
            # be ignored), ==>
            # cannot present final answer yet (LLM may have hallucinated this json)
            return ""
        # valid final answer tool: PASS it on so Critic gets it
        return lr.utils.constants.PASS_TO + "Critic"

    def feedback_tool(self, msg: FeedbackTool) -> str:
        if msg.feedback == "":
            return lr.utils.constants.DONE
        else:
            return f"""
            Below is feedback about your answer. Take it into account to 
            improve your answer, and present it again using the `final_answer_tool`.
            
            FEEDBACK:
            
            {msg.feedback}
            """

    async def llm_response_async(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if self.original_query is None:
            self.original_query = (
                message if isinstance(message, str) else message.content
            )
        result = await super().llm_response_async(message)
        if result is None:
            return result
        # result.content may contain a premature DONE
        # (because weak LLMs tend to repeat their instructions)
        # We deem a DONE to be accidental if no search query results were received
        if not isinstance(message, ChatDocument) or not (
            message.metadata.sender_name == "Searcher"
        ):
            # no search results received yet, so should NOT say DONE
            if isinstance(result, str):
                return result.content.replace(lr.utils.constants.DONE, "")
            result.content = result.content.replace(lr.utils.constants.DONE, "")
            return result

        return result


class CriticAgent(lr.ChatAgent):
    def final_answer_tool(self, msg: FinalAnswerTool) -> str:
        # received from Assistant. Extract the components as plain text,
        # so that the Critic LLM can provide feedback
        return f"""
        The user has presented the following intermediate steps and final answer
        shown below. Please provide feedback using the `feedback_tool`.
        Remember to set the `feedback` field to an empty string if the answer is valid,
        otherwise give specific feedback on what the issues are and how the answer 
        can be improved.
        
        STEPS: {msg.steps}
        
        ANSWER: {msg.answer}
        """

    def feedback_tool(self, msg: FeedbackTool) -> str:
        # say DONE and PASS to the feedback goes back to Assistant to handle
        return lr.utils.constants.DONE + " " + lr.utils.constants.PASS


class SearcherAgentConfig(lr.ChatAgentConfig):
    search_tool_class: Type[lr.ToolMessage]


class SearcherAgent(lr.ChatAgent):
    n_searches: int = 0
    curr_query: str | None = None

    def __init__(self, config: SearcherAgentConfig):
        super().__init__(config)
        self.config: SearcherAgentConfig = config
        self.enable_message(config.search_tool_class)
        self.enable_message(QuestionTool, use=False, handle=True)

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if (
            isinstance(msg, ChatDocument)
            and msg.metadata.sender == lr.Entity.LLM
            and self.n_searches == 0
        ):
            search_tool_name = self.config.search_tool_class.default_value("request")
            return f"""
            You forgot to use the web search tool to answer the 
            user's question : {self.curr_query}.
            REMEMBER - you must ONLY answer the user's questions based on 
             results from a web-search, and you MUST NOT ANSWER them yourself.
             
            Please use the `{search_tool_name}` tool 
            using the specified JSON format, then compose your answer.
            """

    def question_tool(self, msg: QuestionTool) -> str:
        self.curr_query = msg.question
        search_tool_name = self.config.search_tool_class.default_value("request")
        return f"""
        User asked this question: {msg.question}.
        Perform a web search using the `{search_tool_name}` tool
        using the specified JSON format, to find the answer.
        """

    async def llm_response_async(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if (
            isinstance(message, ChatDocument)
            and message.metadata.sender == lr.Entity.AGENT
            and self.n_searches > 0
        ):
            # must be search results from the web search tool,
            # so let the LLM compose a response based on the search results
            self.n_searches = 0  # reset search count

            result = await super().llm_response_forget_async(message)
            # Augment the LLM's composed answer with a helpful nudge
            # back to the Assistant
            result.content = f"""
            Here are the web-search results for the question: {self.curr_query}.
            ===
            {result.content}
            ===
            Decide if you want to ask any further questions, for the 
            user's original question.             
            """
            self.curr_query = None
            return result

        # Handling query from user (or other agent)
        result = await super().llm_response_forget_async(message)
        if result is None:
            return result
        tools = self.get_tool_messages(result)
        if all(not isinstance(t, self.config.search_tool_class) for t in tools):
            # LLM did not use search tool;
            # Replace its response with a placeholder message
            # and the agent fallback_handler will remind the LLM
            result.content = "Did not use web-search tool."
            return result

        self.n_searches += 1
        # result includes a search tool, but may contain DONE in content,
        # so remove that
        result.content = result.content.replace(lr.utils.constants.DONE, "")
        return result


@cl.on_chat_start
async def main(
    debug: bool = True,
    model: str = os.getenv("MODEL", "gpt-4o"),
    nocache: bool = True,
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    await add_instructions(
        title="2-Agent Search Assistant",
        content=dedent(
            """
        Enter a complex question; 
        - The Assistant will break it down into smaller questions for the Searcher
        - The Searcher will search the web and compose a concise answer
        Once the Assistant has enough information, it will say DONE and present the answer.
        
        To answer a new question, click "New Chat".        
        """
        ),
    )

    load_dotenv()

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,
        temperature=0.2,
        max_output_tokens=500,
        timeout=45,
    )

    assistant_config = lr.ChatAgentConfig(
        system_message="""
        You are a resourceful assistant, able to think step by step to answer
        complex questions from the user. You must break down complex questions into
        simpler questions that can be answered by a web search. You must ask me 
        (the user) each question ONE BY ONE, using the `question_tool` in
         the specified format, and I will do a web search and send you
        a brief answer. Once you have enough information to answer my original
        (complex) question, you MUST present your INTERMEDIATE STEPS and FINAL ANSWER
        using the `final_answer_tool` in the specified JSON format.
        You will then receive FEEDBACK from the Critic, and if needed
        you should try to improve your answer based on this feedback.
        """,
        llm=llm_config,
        vecdb=None,
    )
    assistant_agent = AssistantAgent(assistant_config)
    assistant_agent.enable_message(QuestionTool)
    assistant_agent.enable_message(FinalAnswerTool)
    assistant_agent.enable_message(FeedbackTool, use=False, handle=True)

    search_tool_handler_method = MetaphorSearchTool.name()

    search_agent_config = SearcherAgentConfig(
        search_tool_class=MetaphorSearchTool,
        llm=llm_config,
        vecdb=None,
        system_message=f"""
        You are a web-searcher. For ANY question you get, you must use the
        `{search_tool_handler_method}` tool/function-call to get up to 5 results.
        Once you receive the results, you must compose a CONCISE answer 
        based on the search results and say DONE and show the answer to me,
        along with references, in this format:
        DONE [... your CONCISE answer here ...]
        SOURCES: [links from the web-search that you used]
        
        EXTREMELY IMPORTANT: DO NOT MAKE UP ANSWERS, ONLY use the web-search results.
        """,
    )
    search_agent = SearcherAgent(search_agent_config)

    assistant_task = lr.Task(
        assistant_agent,
        name="Assistant",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    search_task = lr.Task(
        search_agent,
        name="Searcher",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )

    critic_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message="""
        You excel at logical reasoning and combining pieces of information.
        The user will send you a summary of the intermediate steps and final answer.
        You must examine these and provide feedback to the user, using the 
        `feedback_tool`, as follows:
        - If you think the answer is valid, 
            simply set the `feedback` field to an empty string "".
        - Otherwise set the `feedback` field to a reason why the answer is invalid,
            and suggest how the user can improve the answer.
        """,
    )
    critic_agent = CriticAgent(critic_agent_config)
    critic_agent.enable_message(FeedbackTool)
    critic_agent.enable_message(FinalAnswerTool, use=False, handle=True)
    critic_task = lr.Task(
        critic_agent,
        name="Critic",
        interactive=False,
    )
    assistant_task.add_sub_task([search_task, critic_task])
    cl.user_session.set("assistant_task", assistant_task)


@cl.on_message
async def on_message(message: cl.Message):
    assistant_task = cl.user_session.get("assistant_task")
    lr.ChainlitTaskCallbacks(assistant_task)
    await assistant_task.run_async(message.content)
</file>

<file path="examples/chainlit/chat-search-assistant.py">
"""
Chainlit version of examples/basic/chat-search-assistant.py,
with only a small change to add the Chainlit callbacks.

See that script for details.

Run like this:

chainlit run examples/chainlit/chat-search-assistant.py

To run with a different LLM, set the MODEL environment variable:

MODEL=ollama/mistral chainlit run examples/chainlit/chat-search-assistant.py

or

MODEL=groq/llama3-70b-8192 chainlit run examples/chainlit/chat-search-assistant.py
"""

import os
from textwrap import dedent

import chainlit as cl
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid.agent.callbacks.chainlit import add_instructions
from langroid.agent.tools.duckduckgo_search_tool import DuckduckgoSearchTool
from langroid.agent.tools.google_search_tool import GoogleSearchTool
from langroid.agent.tools.orchestration import SendTool
from langroid.utils.configuration import Settings, set_global


@cl.on_chat_start
async def main(
    debug: bool = False,
    # e.g. ollama/mistral or local/localhost:5000/v1 default is GPT4o
    model: str = os.getenv("MODEL", ""),
    provider: str = "metaphor",  # or "google", "ddg"
    nocache: bool = False,
):
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    load_dotenv()

    await add_instructions(
        title="2-Agent Search Assistant",
        content=dedent(
            """
        Enter a complex question; 
        - The Assistant will break it down into smaller questions for the Searcher
        - The Searcher will search the web and compose a concise answer
        
        Once the Assistant has enough information, it will say DONE and present the answer.        
        """
        ),
    )

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=8_000,
        temperature=0,
        max_output_tokens=200,
        timeout=45,
    )

    assistant_config = lr.ChatAgentConfig(
        system_message=f"""
        You are a resourceful assistant, able to think step by step to answer
        complex questions from the user. You must break down complex questions into
        simpler questions that can be answered by a web search agent. You must ask 
        each question ONE BY ONE, and the agent will do a web search and send you
        a brief answer. 
        Once you have enough information to answer my original
        (complex) question, you MUST use the TOOL `{SendTool.name()}`
        with `to` set to "User" to send me the answer. 
        """,
        llm=llm_config,
        vecdb=None,
    )
    assistant_agent = lr.ChatAgent(assistant_config)
    assistant_agent.enable_message(SendTool)

    match provider:
        case "google":
            search_tool_class = GoogleSearchTool
        case "metaphor":
            from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool

            search_tool_class = MetaphorSearchTool
        case "ddg":
            search_tool_class = DuckduckgoSearchTool
        case _:
            raise ValueError(f"Unsupported provider {provider} specified.")

    search_tool_handler_method = search_tool_class.default_value("request")

    search_agent_config = lr.ChatAgentConfig(
        llm=llm_config,
        vecdb=None,
        system_message=f"""
        You are a web-searcher. For any question you get, you must use the
        `{search_tool_handler_method}` tool/function-call to get up to 5 results.
        I WILL SEND YOU THE RESULTS; DO NOT MAKE UP THE RESULTS!!
        Once you receive the results, you must compose a CONCISE answer 
        based on the search results and say DONE and show the answer to me,
        in this format:
        DONE [... your CONCISE answer here ...]
        IMPORTANT: YOU MUST WAIT FOR ME TO SEND YOU THE 
        SEARCH RESULTS BEFORE saying you're DONE.
        """,
    )
    search_agent = lr.ChatAgent(search_agent_config)
    search_agent.enable_message(search_tool_class)

    assistant_task = lr.Task(
        assistant_agent,
        name="Assistant",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    search_task = lr.Task(
        search_agent,
        name="Searcher",
        llm_delegate=True,
        single_round=False,
        interactive=False,
    )
    assistant_task.add_sub_task(search_task)
    cl.user_session.set("assistant_task", assistant_task)


@cl.on_message
async def on_message(message: cl.Message):
    assistant_task = cl.user_session.get("assistant_task")
    lr.ChainlitTaskCallbacks(assistant_task)
    await assistant_task.run_async(message.content)
</file>

<file path="examples/chainlit/chat-search-rag.py">
"""
Single-agent question-answering system that has access to
Metaphor web search when needed,
and in case a web search is used, ingests contents into a vector-db,
and uses Retrieval Augmentation to answer the question.

This is a chainlit UI version of examples/docqa/chat-search.py

Run like this:

    chainlit run examples/chainlit/chat-search-rag.py


(See here for guide to using local LLMs with Langroid:)
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import logging
from textwrap import dedent
from typing import Any, List, Optional

import chainlit as cl
import typer

import langroid as lr
import langroid.language_models as lm
from langroid.agent.callbacks.chainlit import (
    add_instructions,
    make_llm_settings_widgets,
    setup_llm,
    update_llm,
)
from langroid.agent.chat_agent import ChatAgent, ChatDocument
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import ForwardTool
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.parsing.web_search import metaphor_search
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

logger = logging.getLogger(__name__)

app = typer.Typer()


class RelevantExtractsTool(ToolMessage):
    request: str = "relevant_extracts"
    purpose: str = (
        "Get docs/extracts relevant to the <query>, from prior search results"
    )
    query: str

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(query="when was the Mistral LLM released?"),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """


class RelevantSearchExtractsTool(ToolMessage):
    request: str = "relevant_search_extracts"
    purpose: str = (
        "Perform an internet search for up to <num_results> results "
        "relevant to the <query>"
    )

    query: str
    num_results: int = 3

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(
                query="when was the Mistral LLM released?",
                num_results=3,
            ),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """


class SearchDocChatAgent(DocChatAgent):
    tried_vecdb: bool = False

    def llm_response_async(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        return ChatAgent.llm_response_async(self, message)

    def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            # non-tool LLM msg => forward to User
            return ForwardTool(agent="User")

    def relevant_extracts(self, msg: RelevantExtractsTool) -> str:
        """Get docs/extracts relevant to the query, from vecdb"""
        self.tried_vecdb = True
        self.callbacks.show_start_response(entity="agent")
        query = msg.query
        logger.info(f"Trying to get relevant extracts for query: {query}")
        _, extracts = self.get_relevant_extracts(query)
        if len(extracts) == 0:
            return """
            No extracts found! You can try doing a web search with the
            `relevant_search_extracts` tool/function-call.
            """
        return "\n".join(str(e) for e in extracts)

    def relevant_search_extracts(self, msg: RelevantSearchExtractsTool) -> str:
        """Get docs/extracts relevant to the query, from a web search"""
        if not self.tried_vecdb and len(self.original_docs) > 0:
            return "Please try the `relevant_extracts` tool, before using this tool"
        query = msg.query
        num_results = msg.num_results
        self.callbacks.show_start_response(entity="agent")
        results = metaphor_search(query, num_results)
        links = [r.link for r in results]
        self.config.doc_paths = links
        self.ingest()
        _, extracts = self.get_relevant_extracts(query)
        if len(extracts) == 0:
            return """
            No release search results found! You can try 
            rephrasing your query to see if results improve, using the
            `relevant_search_extracts` tool/function-call.
            """
        return "\n".join(str(e) for e in extracts)


async def setup_agent_task():
    """Set up Agent and Task from session settings state."""

    # set up LLM and LLMConfig from settings state
    await setup_llm()
    llm_config = cl.user_session.get("llm_config")

    set_global(
        Settings(
            debug=False,
            cache=True,
        )
    )

    config = DocChatAgentConfig(
        name="Searcher",
        llm=llm_config,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        system_message=f"""
        You are a savvy, tenacious, persistent researcher, who knows when to search the 
        internet for an answer.
        
        You will try your best to answer my questions,
        in this order of preference:
        1. If you can answer from your own knowledge, simply return the answer
        2. Otherwise, use the `relevant_extracts` tool/function to
            ask me for some relevant text, and I will send you.  
            Then answer based on the relevant text.
            If I say {NO_ANSWER}, it means I found no relevant docs, and you can try 
            the next step, using a web search.
        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. Answer the question
           based on these text pieces.
        4. If you still can't answer, simply say {NO_ANSWER} 
        5. Be tenacious and persistent, DO NOT GIVE UP. Try asking your questions
        differently to arrive at an 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`.
        
        Be very concise in your responses, use no more than 1-2 sentences.
        When you answer based on provided documents, be sure to show me 
        the SOURCE(s) and EXTRACT(s), for example:
        
        SOURCE: https://www.wikihow.com/Be-a-Good-Assistant-Manager
        EXTRACT: Be a Good Assistant ... requires good leadership skills.
        
        For the EXTRACT, ONLY show up to first 3 words, and last 3 words.
        """,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=200,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    agent = SearchDocChatAgent(config)
    agent.enable_message(RelevantExtractsTool)
    agent.enable_message(RelevantSearchExtractsTool)
    collection_name = "chainlit-chat-search-rag"

    agent.vecdb.set_collection(collection_name, replace=True)

    # set up task with interactive=False, so awaits user ONLY
    # when LLM sends  non-tool msg (see handle_message_fallback method).
    task = Task(agent, interactive=False)
    cl.user_session.set("agent", agent)
    cl.user_session.set("task", task)


@cl.on_settings_update
async def on_update(settings):
    await update_llm(settings)
    await setup_agent_task()


@cl.on_chat_start
async def chat() -> None:
    await add_instructions(
        title="Welcome to the Internet Search + RAG chatbot!",
        content=dedent(
            """
        Ask me anything, especially about recent events that I may not have been trained on.
        
        I have access to two Tools, which I will try to use in order of priority:
        - `relevant_extracts` to try to answer your question using Retrieval Augmented Generation
           from prior search results ingested into a vector-DB (from prior searches in this session),
           and failing this, I will use my second tool:
        - `relevant_search_extracts` to do a web search (Using Metaphor Search)
        and ingest the results into the vector-DB, and then use 
        Retrieval Augmentation Generation (RAG) to answer the question.
        """
        ),
    )

    await make_llm_settings_widgets(
        lm.OpenAIGPTConfig(
            timeout=180,
            chat_context_length=16_000,
            chat_model="",
            temperature=0.1,
        )
    )
    await setup_agent_task()


@cl.on_message
async def on_message(message: cl.Message):
    task = cl.user_session.get("task")
    lr.ChainlitTaskCallbacks(task)
    await task.run_async(message.content)
</file>

<file path="examples/chainlit/chat-search.py">
"""
Basic single-agent chat example, using a web Search Tool, using ChainlitTaskCallbacks.

- User asks a question
- LLM either responds directly or generates a Metaphor web search Tool/function-call
    - if Tool used:
         - Agent handler recognizes this tool and returns search results
         - LLM sees search results and composes a response.
- user asks another question


After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-search.py
"""

import logging
from textwrap import dedent
from typing import Optional

import chainlit as cl

import langroid as lr
from langroid import ChatDocument
from langroid.agent.callbacks.chainlit import (
    add_instructions,
    make_llm_settings_widgets,
    setup_llm,
    update_llm,
)
from langroid.agent.tools.duckduckgo_search_tool import DuckduckgoSearchTool
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool

logger = logging.getLogger(__name__)


def search_system_message(search_tool: lr.ToolMessage) -> str:
    tool_name = search_tool.default_value("request")
    sys_msg = f"""
        You are an astute, self-aware AI assistant, and you are adept at 
        responding to a user's question in one of two ways:
        - If you KNOW the answer from your own knowledge, respond directly.
        - OTHERWISE, request up to 5 results from a web search using 
          the `{tool_name}` tool/function-call.
          In this case you will receive the web search results, and you can 
          then compose a response to the user's question. 
    """
    return sys_msg


class SearchAgent(lr.ChatAgent):
    async def user_response_async(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        response = await super().user_response_async(message)
        if response is None:
            return None
        content = response.content
        search_tool = MetaphorSearchTool
        if content.startswith("/"):
            match content[1]:
                case "d":
                    search_tool = DuckduckgoSearchTool
                    self.enable_message(DuckduckgoSearchTool)
                    self.enable_message(MetaphorSearchTool, use=False, handle=False)
                case "m":
                    search_tool = MetaphorSearchTool
                    self.enable_message(MetaphorSearchTool)
                    self.enable_message(DuckduckgoSearchTool, use=False, handle=False)

            self.clear_history(0)
            sys_msg = search_system_message(search_tool)
            self.set_system_message(sys_msg)

            response.content = content[2:]
        return response

    async def agent_response_async(self, message: ChatDocument) -> ChatDocument:
        response = await super().agent_response_async(message)
        if response is None:
            return None
        # ensure tool result goes to LLM
        response.metadata.recipient = lr.Entity.LLM
        return response


async def setup_agent_task(search_tool: lr.ToolMessage):
    """Set up Agent and Task from session settings state."""

    # set up LLM and LLMConfig from settings state
    await setup_llm()
    llm_config = cl.user_session.get("llm_config")
    sys_msg = search_system_message(search_tool)
    config = lr.ChatAgentConfig(
        llm=llm_config,
        name="Searcher",
        system_message=sys_msg,
    )
    agent = SearchAgent(config)
    agent.enable_message(search_tool)
    task = lr.Task(agent, interactive=True)
    cl.user_session.set("agent", agent)
    cl.user_session.set("task", task)


@cl.on_settings_update
async def on_update(settings):
    await update_llm(settings)
    await setup_agent_task(MetaphorSearchTool)


@cl.on_chat_start
async def on_chat_start():
    await add_instructions(
        title="Agent with access to a web search Tool",
        content=dedent(
            """
        Agent uses a tool/fn-call to search the web 
        
        Default search is using DuckDuckGo. You can switch the search to 
        - Duckduckgo by typing `/d` at the start of your question
        - Metaphor by typing `/m` at the start of your question
        
        This is the flow:
        - User asks question
        - Agent LLM uses an internet search tool to generate search results
        - Agent handler recognizes this tool and returns search results
        - User hits `c` to continue
        - Agent LLM composes answer
        
        To change LLM settings, including model name, click the settings symbol on the 
        left of the chat window.        
        """
        ),
    )

    await make_llm_settings_widgets()
    await setup_agent_task(MetaphorSearchTool)


@cl.on_message
async def on_message(message: cl.Message):
    task = cl.user_session.get("task")
    lr.ChainlitTaskCallbacks(task)
    await task.run_async(message.content)
</file>

<file path="examples/chainlit/chat-tool.py">
"""
Basic single-agent chat example, using task.run(), with a tool, with streaming,
using ChainlitTaskCallbacks.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-tool.py
"""

from textwrap import dedent

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import add_instructions


class CapitalTool(lr.ToolMessage):
    request: str = "capital"
    purpose: str = "To present the capital of given <country>."
    country: str
    capital: str

    def handle(self) -> str:
        return f"""
        Success! LLM responded with a tool/function-call, with result:
        
        Capital of {self.country} is {self.capital}.
        """


@cl.on_chat_start
async def on_chat_start():
    config = lr.ChatAgentConfig(
        name="CapitalExpert",
        system_message="""
        When asked for the <capital> of a <country>, present
        your response using the `capital` tool/function-call.
        """,
    )
    agent = lr.ChatAgent(config)
    agent.enable_message(CapitalTool)

    await add_instructions(
        title="Instructions",
        content=dedent(
            """
        Interact with a **Langroid Task**, whose ChatAgent has access 
        to a `capital` tool. You can ask about anything, but whenever you ask 
        about a country's capital, the agent will use the `capital` tool to present 
        the capital of that country. This "tool-message" is handled by the Agent's 
        handler method, and the result is presented as plain text.
        """
        ),
    )
    # inject callbacks into the agent
    task = lr.Task(
        agent,
        interactive=True,
    )
    cl.user_session.set("task", task)


@cl.on_message
async def on_message(message: cl.Message):
    task = cl.user_session.get("task")
    lr.ChainlitTaskCallbacks(task)
    await task.run_async(message.content)
</file>

<file path="examples/chainlit/chat-transcript.py">
"""
Variant of chat-agent.py, that waits for user to type "/s" (meaning submit)
to store chat transcript in a file.

Directly uses an Agent (i.e. without Task) 
using callbacks, which also enables streaming.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-transcript.py

or:
    
uv run chainlit run examples/chainlit/chat-transcript.py

"""

import logging

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import add_instructions

# set info logger
logging.basicConfig(level=logging.INFO)

FILE = "examples/chainlit/chat-transcript.txt"


@cl.on_chat_start
async def on_chat_start():
    config = lr.ChatAgentConfig(
        name="Demo",
        system_message="You are a helpful assistant. Be concise in your answers.",
    )
    agent = lr.ChatAgent(config)

    cl.user_session.set("agent", agent)

    await add_instructions(
        title="Instructions",
        content="Interact with a **Langroid ChatAgent**",
    )


@cl.on_message
async def on_message(message: cl.Message):
    agent: lr.ChatAgent = cl.user_session.get("agent")
    # important: only apply callbacks after getting first msg.
    lr.ChainlitAgentCallbacks(agent)
    if message.content.startswith("/s"):
        content = message.content
        # get transcript of entire conv history as a string
        history = (
            "\n\n".join(
                [
                    f"{msg.role.value.upper()}: {msg.content}"
                    for msg in agent.message_history
                ]
            )
            + "\n\n"
            + "FINAL User Answer: "
            + content[2:]
        )

        # save chat transcript to file
        with open(FILE, "w") as f:
            f.write(f"Chat transcript:\n\n{history}\n")
            await cl.Message(
                content=f"Chat transcript saved to {FILE}.",
                author="System",
            ).send()
        return

    await agent.llm_response_async(message.content)
</file>

<file path="examples/chainlit/chat-tree-chainlit.py">
"""
Variant of chat-tree.py but with Chainlit UI.
The ONLY change is we apply ChainlitTaskCallbacks() to the top-level task!

Run like this:

chainlit run examples/chainlit/chat-tree-chainlit.py
"""

from textwrap import dedent

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import add_instructions
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE
from langroid.utils.globals import GlobalState

INTERACTIVE = False


class MyGlobalState(GlobalState):
    number: int | None = None


class AskNumTool(ToolMessage):
    request: str = "ask_num"
    purpose: str = "Ask user for the initial number"


class AddNumTool(ToolMessage):
    request: str = "add_num"
    purpose: str = "Add <number> to the original number, return the result"
    number: int

    def handle(self) -> str:
        """
        This is a stateless tool (i.e. does not use any Agent member vars), so we can
        define the handler right here, instead of defining an `add_num`
        method in the agent.
        """
        return str(int(MyGlobalState.get_value("number")) + int(self.number))


class MainChatAgent(ChatAgent):
    def ask_num(self, msg: AskNumTool) -> str:
        res = self.callbacks.get_user_response(prompt="Please enter a number")
        # record this in global state, so other agents can access it
        num = res
        MyGlobalState.set_values(number=num)
        return str(num)


@cl.on_chat_start
async def on_start():
    await add_instructions(
        title="Multi-agent chat for tree-structured computation with tools",
        content=dedent(
            """
        This task consists of performing this calculation for a given input number n:
        
        ```python
        def Main(n):
            if n is odd:
                return (3*n+1) + n
            else:
                If n is divisible by 10:
                    return n/10 + n
                else:
                    return n/2 + n
        ```
        
        See details in the [chat-tree.py](https://github.com/langroid/langroid/blob/main/examples/basic/chat-tree.py), 
        and the writeup on 
        [Hierarchical Agent Computation](https://langroid.github.io/langroid/examples/agent-tree/).
        
        To start the computation, enter a number.  
        """
        ),
    )


@cl.on_message
async def chat(msg: cl.Message) -> None:
    set_global(
        Settings(
            debug=False,
            cache=True,
            stream=True,
        )
    )

    MyGlobalState.set_values(number=int(msg.content))

    config = ChatAgentConfig(
        llm=OpenAIGPTConfig(
            chat_model=OpenAIChatModel.GPT4o,
        ),
        vecdb=None,
    )

    main_agent = MainChatAgent(config)
    main_task = Task(
        main_agent,
        name="Main",
        interactive=INTERACTIVE,
        system_message="""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        In this case simply write the <number>, say nothing else.
        
        RESULT Message format: RESULT <number>
        In this case simply say "DONE <number>", e.g.:
        DONE 19
        """,
    )

    # Handles only even numbers
    even_agent = ChatAgent(config)
    even_task = Task(
        even_agent,
        name="Even",
        interactive=INTERACTIVE,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if the <number> is odd, say '{DONE}'
        - otherwise, simply write the <number>, say nothing else.
        
        RESULT Message format: RESULT <number>
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    # handles only even numbers ending in Zero
    evenz_agent = ChatAgent(config)
    evenz_task = Task(
        evenz_agent,
        name="EvenZ",
        interactive=INTERACTIVE,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if <number> n is even AND divisible by 10, compute n/10 and pass it on,
        - otherwise, say '{DONE}'
        
        RESULT Message format: RESULT <number>
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    # Handles only even numbers NOT ending in Zero
    even_nz_agent = ChatAgent(config)
    even_nz_task = Task(
        even_nz_agent,
        name="EvenNZ",
        interactive=INTERACTIVE,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if <number> n is even AND NOT divisible by 10, compute n/2 and pass it on,
        - otherwise, say '{DONE}'
        
        RESULT Message format: RESULT <number>
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    # Handles only odd numbers
    odd_agent = ChatAgent(config)
    odd_task = Task(
        odd_agent,
        name="Odd",
        interactive=INTERACTIVE,
        system_message=f"""
        You will receive two types of messages, to which you will respond as follows:
        
        INPUT Message format: <number>
        - if <number> n is odd, compute n*3+1 and write it.
        - otherwise, say '{DONE}'

        RESULT Message format: RESULT <number>        
        In this case simply write "DONE RESULT <number>", e.g.:
        DONE RESULT 19
        """,
    )

    adder_agent = ChatAgent(config)
    # set up the tools
    adder_agent.enable_message(AddNumTool)
    # main_agent.enable_message(AskNumTool)

    adder_task = Task(
        adder_agent,
        name="Adder",
        interactive=INTERACTIVE,
        system_message="""
        You will be given a number n.
        You have to add it to the original number and return the result.
        You do not know the original number, so you must use the 
        `add_num` tool/function for this. 
        When you receive the result, say "DONE RESULT <result>", e.g.
        DONE RESULT 19
        """,
    )

    # set up tasks and subtasks
    main_task.add_sub_task([even_task, odd_task])
    even_task.add_sub_task([evenz_task, even_nz_task])
    evenz_task.add_sub_task(adder_task)
    even_nz_task.add_sub_task(adder_task)
    odd_task.add_sub_task(adder_task)

    # inject chainlit callbacks: this is the ONLY change to chat-tree.py
    lr.ChainlitTaskCallbacks(main_task)

    # start the chat
    await main_task.run_async(msg.content)
</file>

<file path="examples/chainlit/chat-with-agent.py">
"""
Basic single-agent chat example, to directly use an Agent (i.e. without Task)
using callbacks, which also enables streaming.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-with-agent.py

"""

import logging

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import add_instructions

# set info logger
logging.basicConfig(level=logging.INFO)


@cl.on_chat_start
async def on_chat_start():
    config = lr.ChatAgentConfig(
        name="Demo",
        system_message="You are a helpful assistant. Be concise in your answers.",
    )
    agent = lr.ChatAgent(config)

    cl.user_session.set("agent", agent)

    await add_instructions(
        title="Instructions",
        content="Interact with a **Langroid ChatAgent**",
    )


@cl.on_message
async def on_message(message: cl.Message):
    agent: lr.ChatAgent = cl.user_session.get("agent")
    # important: only apply callbacks after getting first msg.
    lr.ChainlitAgentCallbacks(agent)
    await agent.llm_response_async(message.content)
</file>

<file path="examples/chainlit/chat-with-task.py">
"""
Basic single-agent chat example using Task along with ChainlitTaskCallbacks.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/chat-with-task.py
"""

from textwrap import dedent

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import (
    add_instructions,
    make_llm_settings_widgets,
    setup_llm,
    update_llm,
)


@cl.on_settings_update
async def on_settings_update(settings: cl.ChatSettings):
    await update_llm(settings)
    await setup_agent_task()


async def setup_agent_task():
    await setup_llm()
    llm_config = cl.user_session.get("llm_config")

    config = lr.ChatAgentConfig(
        llm=llm_config,
        name="Demo",
        system_message="You are a helpful assistant. Be concise in your answers.",
    )
    agent = lr.ChatAgent(config)

    task = lr.Task(
        agent,
        interactive=True,
    )
    cl.user_session.set("task", task)


@cl.on_chat_start
async def on_chat_start():
    await add_instructions(
        title="Basic Langroid Chatbot",
        content=dedent(
            """
        Uses Langroid's `Task.run()`. 
        Before starting the chat, 
        you can change LLM settings by clicking the settings icon next to the chat window.
        """
        ),
    )
    await make_llm_settings_widgets()
    await setup_agent_task()


@cl.on_message
async def on_message(message: cl.Message):
    task = cl.user_session.get("task")
    # sometimes we may want the User to NOT have agent name in front,
    # and just show them as YOU.
    callback_config = lr.ChainlitCallbackConfig(user_has_agent_name=False)
    lr.ChainlitTaskCallbacks(task, config=callback_config)
    await task.run_async(message.content)
</file>

<file path="examples/chainlit/cypher_message.py">
CONSTRUCT_DEPENDENCY_GRAPH = """
        with "{package_type}" as system, "{package_name}" as name, "{package_version}" as version

        call apoc.load.model_dump_json("https://api.deps.dev/v3alpha/systems/"+system+"/packages/"
                            +name+"/versions/"+version+":dependencies")
        yield value as r
        
        call {{ with r
                unwind r.nodes as package
                merge (p:Package:PyPi {{name: package.versionKey.name, version: package.versionKey.version}})
                return collect(p) as packages
        }}
        call {{ with r, packages
            unwind r.edges as edge
            with packages[edge.fromNode] as from, packages[edge.toNode] as to, edge
            merge (from)-[rel:DEPENDS_ON]->(to) ON CREATE SET rel.requirement 
            = edge.requirement
            return count(*) as numRels
        }}
        
        match (root:Package:PyPi) where root.imported is null
        set root.imported = true
        with "{package_type}" as system, root.name as name, root.version as version
        call apoc.load.model_dump_json("https://api.deps.dev/v3alpha/systems/"+system+"/packages/"
                            +name+"/versions/"+version+":dependencies")
        yield value as r
        
        call {{ with r
                unwind r.nodes as package
                merge (p:Package:PyPi {{name: package.versionKey.name, version: package.versionKey.version}})
                return collect(p) as packages
        }}
        call {{ with r, packages
                unwind r.edges as edge
                with packages[edge.fromNode] as from, packages[edge.toNode] as to, edge
                merge (from)-[rel:DEPENDS_ON]->(to) ON CREATE SET 
                rel.requirement = edge.requirement
                return count(*) as numRels
        }}
        return size(packages) as numPackages, numRels
        """
</file>

<file path="examples/chainlit/dependency_chatbot.py">
"""
Single-agent to use to chat with a Neo4j knowledge-graph (KG)
that models a dependency graph of Python packages.

This is a chainlit UI version of examples/kg-chat/dependency_chatbot.py

Run like this:
```
chainlit run examples/kg-chat/dependency_chatbot.py
```

The requirements are described in
 `https://github.com/langroid/langroid/blob/main/examples/kg-chat/README.md`
"""

import webbrowser
from pathlib import Path
from textwrap import dedent

import chainlit as cl
import typer
from cypher_message import CONSTRUCT_DEPENDENCY_GRAPH
from pyvis.network import Network
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.agent.callbacks.chainlit import (
    add_instructions,
    make_llm_settings_widgets,
    setup_llm,
    update_llm,
)
from langroid.agent.special.neo4j.neo4j_chat_agent import (
    Neo4jChatAgent,
    Neo4jChatAgentConfig,
    Neo4jSettings,
)
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.google_search_tool import GoogleSearchTool
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

app = typer.Typer()


class DepGraphTool(ToolMessage):
    request: str = "construct_dependency_graph"
    purpose: str = f"""Get package <package_version>, <package_type>, and <package_name>.
    For the <package_version>, obtain the recent version, it should be a number. 
    For the <package_type>, return if the package is PyPI or not.
      Otherwise, return {NO_ANSWER}.
    For the <package_name>, return the package name provided by the user.
    ALL strings are in lower case.
    """
    package_version: str
    package_type: str
    package_name: str


class VisualizeGraph(ToolMessage):
    request: str = "visualize_dependency_graph"
    purpose: str = """
      Use this tool/function to display the dependency graph.
      """
    package_version: str
    package_type: str
    package_name: str
    query: str


class DependencyGraphAgent(Neo4jChatAgent):
    def construct_dependency_graph(self, msg: DepGraphTool) -> None:
        check_db_exist = (
            "MATCH (n) WHERE n.name = $name AND n.version = $version RETURN n LIMIT 1"
        )
        response = self.read_query(
            check_db_exist, {"name": msg.package_name, "version": msg.package_version}
        )
        if response.success and response.data:
            # self.config.database_created = True
            return "Database Exists"
        else:
            construct_dependency_graph = CONSTRUCT_DEPENDENCY_GRAPH.format(
                package_type=msg.package_type.lower(),
                package_name=msg.package_name,
                package_version=msg.package_version,
            )
            response = self.write_query(construct_dependency_graph)
            if response.success:
                self.config.database_created = True
                return "Database is created!"
            else:
                return f"""
                    Database is not created!
                    Seems the package {msg.package_name} is not found,
                    """

    def visualize_dependency_graph(self, msg: VisualizeGraph) -> str:
        """
        Visualizes the dependency graph based on the provided message.

        Args:
            msg (VisualizeGraph): The message containing the package info.

        Returns:
            str: response indicates whether the graph is displayed.
        """
        # Query to fetch nodes and relationships
        # TODO: make this function more general to return customized graphs
        # i.e, displays paths or subgraphs
        query = """
            MATCH (n)
            OPTIONAL MATCH (n)-[r]->(m)
            RETURN n, r, m
        """

        query_result = self.read_query(query)
        nt = Network(notebook=False, height="750px", width="100%", directed=True)

        node_set = set()  # To keep track of added nodes

        for record in query_result.data:
            # Process node 'n'
            if "n" in record and record["n"] is not None:
                node = record["n"]
                # node_id = node.get("id", None)  # Assuming each node has a unique 'id'
                node_label = node.get("name", "Unknown Node")
                node_title = f"Version: {node.get('version', 'N/A')}"
                node_color = "blue" if node.get("imported", False) else "green"

                # Check if node has been added before
                if node_label not in node_set:
                    nt.add_node(
                        node_label, label=node_label, title=node_title, color=node_color
                    )
                    node_set.add(node_label)

            # Process relationships and node 'm'
            if (
                "r" in record
                and record["r"] is not None
                and "m" in record
                and record["m"] is not None
            ):
                source = record["n"]
                target = record["m"]
                relationship = record["r"]

                source_label = source.get("name", "Unknown Node")
                target_label = target.get("name", "Unknown Node")
                relationship_label = (
                    relationship[1]
                    if isinstance(relationship, tuple) and len(relationship) > 1
                    else "Unknown Relationship"
                )

                # Ensure both source and target nodes are added before adding the edge
                if source_label not in node_set:
                    source_title = f"Version: {source.get('version', 'N/A')}"
                    source_color = "blue" if source.get("imported", False) else "green"
                    nt.add_node(
                        source_label,
                        label=source_label,
                        title=source_title,
                        color=source_color,
                    )
                    node_set.add(source_label)
                if target_label not in node_set:
                    target_title = f"Version: {target.get('version', 'N/A')}"
                    target_color = "blue" if target.get("imported", False) else "green"
                    nt.add_node(
                        target_label,
                        label=target_label,
                        title=target_title,
                        color=target_color,
                    )
                    node_set.add(target_label)

                nt.add_edge(source_label, target_label, title=relationship_label)

        nt.options.edges.font = {"size": 12, "align": "top"}
        nt.options.physics.enabled = True
        nt.show_buttons(filter_=["physics"])

        output_file_path = "neo4j_graph.html"
        nt.write_html(output_file_path)

        # Try to open the HTML file in a browser
        try:
            abs_file_path = str(Path(output_file_path).resolve())
            webbrowser.open("file://" + abs_file_path, new=2)
        except Exception as e:
            print(f"Failed to automatically open the graph in a browser: {e}")


async def setup_agent_task():
    """Set up Agent and Task from session settings state."""

    # set up LLM and LLMConfig from settings state
    await setup_llm()
    llm_config = cl.user_session.get("llm_config")

    set_global(
        Settings(
            debug=False,
            cache=True,
        )
    )

    neo4j_settings = Neo4jSettings()

    dependency_agent = DependencyGraphAgent(
        config=Neo4jChatAgentConfig(
            neo4j_settings=neo4j_settings,
            show_stats=False,
            llm=llm_config,
        ),
    )

    system_message = f"""You are an expert in Dependency graphs and analyzing them using
    Neo4j. 
    
    FIRST, I'll give you the name of the package that I want to analyze.
    
    THEN, you can also use the `web_search` tool/function to find out information about a package,
      such as version number and package type (PyPi or not). 
    
    If unable to get this info, you can ask me and I can tell you.
    
    DON'T forget to include the package name in your questions. 
      
    After receiving this information, make sure the package version is a number and the
    package type is PyPi.
    THEN ask the user if they want to construct the dependency graph,
    and if so, use the tool/function `construct_dependency_graph` to construct
      the dependency graph. Otherwise, say `Couldn't retrieve package type or version`
      and {NO_ANSWER}.
    After constructing the dependency graph successfully, you will have access to Neo4j 
    graph database, which contains dependency graph.
    You will try your best to answer my questions. Note that:
    1. You can use the tool `get_schema` to get node label and relationships in the
    dependency graph. 
    2. You can use the tool `retrieval_query` to get relevant information from the
      graph database. I will execute this query and send you back the result.
      Make sure your queries comply with the database schema.
    3. Use the `web_search` tool/function to get information if needed.
    To display the dependency graph use this tool `visualize_dependency_graph`.
    """
    task = Task(
        dependency_agent,
        name="DependencyAgent",
        system_message=system_message,
    )

    dependency_agent.enable_message(DepGraphTool)
    dependency_agent.enable_message(GoogleSearchTool)
    dependency_agent.enable_message(VisualizeGraph)

    cl.user_session.set("dependency_agent", dependency_agent)
    cl.user_session.set("task", task)


@cl.on_settings_update
async def on_update(settings):
    await update_llm(settings)
    await setup_agent_task()


@cl.on_chat_start
async def chat() -> None:
    await add_instructions(
        title="Welcome to Python Dependency chatbot!",
        content=dedent(
            """
        Ask any questions about Python packages, and I will try my best to answer them.
        But first, the user specifies package name
        -> agent gets version number and type of package using google search
        -> agent builds dependency graph using Neo4j
        -> user asks natural language query about dependencies
        -> LLM translates to Cypher query to get info from KG
        -> Query results returned to LLM
        -> LLM translates to natural language response
        """
        ),
    )

    await make_llm_settings_widgets(
        lm.OpenAIGPTConfig(
            timeout=180,
            chat_context_length=16_000,
            chat_model="",
            temperature=0.1,
        )
    )
    await setup_agent_task()


@cl.on_message
async def on_message(message: cl.Message):
    task = cl.user_session.get("task")
    lr.ChainlitTaskCallbacks(task)
    await task.run_async(message.content)
</file>

<file path="examples/chainlit/extract-then-chat.py">
"""
3-Agent system to first extract a few pieces of info, then chat with user.

- Assistant: helps user answer questions about a Book. But first it needs to
    extract some information from a document about the Book, using Extractor.
- Extractor: generates questions about the Book document, one by one,
    then returns all info to Assistant using a tool message.
- DocAgent: answers the questions generated by Extractor, based on the Book doc.

Run like this:

chainlit run examples/chainlit/extract-then-chat.py

"""

import os
from textwrap import dedent
from typing import List

import chainlit as cl
from dotenv import load_dotenv

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.callbacks.chainlit import add_instructions
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE, NO_ANSWER, PASS, SEND_TO

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class BookInfo(BaseModel):
    title: str
    author: str
    year: int


class BookInfoTool(ToolMessage):
    request: str = "book_info"
    purpose: str = "Collect <info> about Books"

    info: List[BookInfo]

    def handle(self) -> str:
        """Exit task and pass tool to parent"""
        return DONE + " " + PASS

    @classmethod
    def examples(cls) -> List["BookInfoTool"]:
        return [
            cls(
                info=[
                    BookInfo(title="The Hobbit", author="J.R.R. Tolkien", year=1937),
                    BookInfo(
                        title="The Great Gatsby",
                        author="F. Scott Fitzgerald",
                        year=1925,
                    ),
                ]
            )
        ]


class Assistant(ChatAgent):
    def book_info(self, msg: BookInfoTool) -> str:
        # convert info  to NON-JSON so it doesn't look like a tool,
        # and insert routing so that the Assistan't LLM responds to it, not user.
        info_str = str(msg.info).replace("{", "[").replace("}", "]")
        return f"""{SEND_TO}LLM
        Below is INFO about various books, you received from the Extractor.
        Now ask the user what help they need, and respond ONLY based on this INFO.
        
        INFO: 
        {info_str} 
        """


class Extractor(ChatAgent):
    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        """Nudge LLM when it fails to use book_info correctly"""
        if self.has_tool_message_attempt(msg):
            return """
            You must use the "book_info" tool to present the info.
            You either forgot to use it, or you used it with the wrong format.
            Make sure all fields are filled out and pay attention to the 
            required types of the fields.
            """


@cl.on_chat_start
async def on_chat_start():
    await add_instructions(
        title="Hello! I am your book info helper. "
        "First I will get info about some books",
        content=dedent(
            """
        Enter `x` or `q` to quit at any point.
        """
        ),
    )

    load_dotenv()

    set_global(
        Settings(
            debug=False,
            cache=True,  # disables cache lookup; set to True to use cache
        )
    )

    llm_cfg = lm.OpenAIGPTConfig(
        # or, e.g. "ollama/mistral:7b-instruct-v0.2-q8_0" but result may be brittle
        chat_model=lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust based on model
    )
    doc_agent = DocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            n_neighbor_chunks=2,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=50,
                overlap=10,
                n_neighbor_ids=4,
            ),
            vecdb=lr.vector_store.QdrantDBConfig(
                collection_name="book_info",
                replace_collection=True,
                storage_path=".qdrant/data/",
                cloud=False,
                embedding=lr.embedding_models.SentenceTransformerEmbeddingsConfig(
                    model_type="sentence-transformer",
                    model_name="BAAI/bge-large-en-v1.5",
                ),
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.ingest_doc_paths(["examples/chainlit/books.txt"])
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        # Don't use system_message here since it will override doc chat agent's
        # default system message
    )

    extractor_agent = Extractor(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    extractor_agent.enable_message(BookInfoTool)

    extractor_task = Task(
        extractor_agent,
        name="Extractor",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message=f"""
        You are an expert at understanding JSON function/tool specifications.
        You must extract information about various books from a document,
        to finally present the info using the `book_info` tool/function,
        but you do not have access to the document. 
        I can help with your questions about the document.
        You have to ask questions in these steps:
        1. ask which books are in the document
        2. for each book, ask the various pieces of info you need.
        
        If I am unable to answer your question initially, try asking differently,
        and if I am still unable to answer after 3 tries, 
        fill in {NO_ANSWER} for that field. 
        Think step by step. 
        
        Do not explain yourself, or say any extraneous things. 
        When you receive the answer, then ask for the next field, and so on.
        """,
    )

    assistant_agent = Assistant(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    assistant_agent.enable_message(lr.agent.tools.RecipientTool)
    # enable assistant to HANDLE the book_info tool but not USE it
    assistant_agent.enable_message(BookInfoTool, use=False, handle=True)
    assistant_task = Task(
        assistant_agent,
        name="Assistant",
        interactive=True,
        system_message="""
        You are a helpful librarian, answering my (the user) questions about 
        books described in a certain document, and you do NOT know which 
        books are in the document.
        
        FIRST you need to ask the "Extractor" to collect information
        about various books that are in a certain document. Address your request to the 
        Extractor using the 'recipient_message' tool/function. 
        
        Once you receive the information, you should then ask me (the user) 
        what I need help with.                
        """,
    )

    assistant_task.add_sub_task([extractor_task])
    extractor_task.add_sub_task([doc_task])

    lr.ChainlitTaskCallbacks(assistant_task)
    # must use run() instead of run_async() because DocChatAgent
    # does not have an async llm_response method
    await assistant_task.run_async()
</file>

<file path="examples/chainlit/multi-agent-nested-tool.py">
"""
TODO - this example does not work yet due to breaking changes in Chainlit

2-agent chat, using task.run(), where the sub-task uses a tool to get user input.
This illustrates how a sub-task's steps, including tool-calls, are nested
one level under the parent task's steps.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/multi-agent-nested-tool.py
"""

from textwrap import dedent

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import ChainlitTaskCallbacks, add_instructions
from langroid.utils.configuration import settings
from langroid.utils.constants import DONE

settings.cache = False


class ExportTool(lr.ToolMessage):
    request: str = "main_export"
    purpose: str = "To request the main export of a given <country>."
    country: str


class StudentChatAgent(lr.ChatAgent):
    def main_export(self, msg: ExportTool) -> str:
        assert (
            self.callbacks.get_user_response is not None
        ), "No get_user_response method"
        assert (
            self.callbacks.show_agent_response is not None
        ), "No show_agent_response method"

        prompt = "Please tell me the main export of " + msg.country
        # create the question for user as an agent response since it
        # will ensure it is shown at right nesting level
        # self.callbacks.show_agent_response(content=prompt)
        user_response = self.callbacks.get_user_response(prompt=prompt)
        res = "the main export is " + user_response
        return res


@cl.on_chat_start
async def on_chat_start():
    await add_instructions(
        title="Two-Agent Demo, where sub-agent uses a Tool/function-call",
        content=dedent(
            """
        **Teacher Agent** delegates to **Student Agent.** 
        - **Teacher** Agent asks a "country export" question to **Student** Agent
        - user (you) hits `c` to continue on to the **Student**
        - **Student** LLM uses `export` tool/fn-call to get answer from user
        - **Student** Agent handler code presents this question to you (user)
        - you answer the question
        - **Student** Agent handler returns your answer
        - **Student** LLM shows the answer
        - user hits `c` to continue on to the **Teacher**
        - **Teacher** Agent gives feedback
        - and so on.
        
        Note how all steps of the (student) sub-task are nested one level below 
        the main (teacher) task.
        """
        ),
    )

    config = lr.ChatAgentConfig()
    teacher_agent = lr.ChatAgent(config)
    teacher_task = lr.Task(
        teacher_agent,
        name="Teacher",
        interactive=True,
        system_message="""
        Ask your student what the main export of a country is, and give feedback. 
        Start with a question!
        """,
    )
    student_agent = StudentChatAgent(config)
    student_agent.enable_message(ExportTool)
    student_task = lr.Task(
        student_agent,
        name="Student",
        interactive=True,
        system_message=f"""
        When you receive a country-export question, 
        use the `main_export` tool to get the answer from the user.
        When you get the answer, say {DONE} and show the answer.
        """,
    )

    teacher_task.add_sub_task(student_task)
    ChainlitTaskCallbacks(teacher_task)
    await teacher_task.run_async()
</file>

<file path="examples/chainlit/multi-agent.py">
"""
2-agent chat, using task.run(), where:
- Teacher Agent asks a question
- Student Agent answers the question
- Teacher Agent gives feedback
- ...


After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/multi-agent.py
"""

import os
from textwrap import dedent

import chainlit as cl

import langroid as lr
from langroid.agent.callbacks.chainlit import ChainlitTaskCallbacks, add_instructions
from langroid.utils.configuration import settings


@cl.on_chat_start
async def on_chat_start(
    debug: bool = os.getenv("DEBUG", False),
    no_cache: bool = os.getenv("NOCACHE", False),
):
    settings.debug = debug
    settings.cache = not no_cache

    await add_instructions(
        title="Two-Agent Demo",
        content=dedent(
            """
        **Teacher Agent** delegates to **Student Agent.**
        - **Teacher** Agent asks a numerical question to **Student** Agent
        - **Student** Agent answers the question
        - **Teacher** Agent gives feedback        
        - and so on until 10 turns are done.
        
        Note how all steps of the (student) sub-task are nested one level below 
        the main (teacher) task.
        """
        ),
    )
    config = lr.ChatAgentConfig()
    teacher_agent = lr.ChatAgent(config)
    teacher_task = lr.Task(
        teacher_agent,
        name="Teacher",
        interactive=False,
        system_message="""
        Ask your student concise numerical questions, and give feedback. 
        Start with a question!
        """,
    )
    student_agent = lr.ChatAgent(config)
    student_task = lr.Task(
        student_agent,
        name="Student",
        interactive=False,
        system_message="""Concisely answer your teacher's numerical questions""",
        single_round=True,
    )

    teacher_task.add_sub_task(student_task)
    ChainlitTaskCallbacks(teacher_task)
    await teacher_task.run_async(turns=10)
</file>

<file path="examples/chainlit/multi-extract-3.py">
"""
TODO: Fix this example, it fails due to breaking changes in Chainlit

3-Agent system to extract structured information from a document.
(This is a chainlit version of examples/docqa/chat-multi-extract-3.py)

- LeaseExtractor: is tasked with extracting structured information from a
    commercial lease document, and must present the terms in a specific nested JSON
    format. This agent generates questions corresponding to each field in the JSON
    format.
- Validator: This agent detects if LeaseExtractorAgent's message is asking for ONE
    piece of information, or MULTIPLE pieces. If the message is only asking about ONE
    thing, OR if it is NOT EVEN a question, it responds with "DONE" and says nothing.
    If the message is asking MORE THAN ONE thing, it responds with a message asking to
    only ask ONE question at a time.
    [Why restrict to one question at a time? Because the DocAgent is more likely to
    understand and answer a single question at a time]

- DocAgent: This agent answers the questions generated by LeaseExtractorAgent,
    based on the lease document it has access to via vecdb, using RAG.

Run like this:

```
chainlit run examples/chainlit/multi-extract-3.py
```

Edit the `model` argument in main() fn below to change the model.
If you set it to "", it will default to the GPT4-turbo model.


For more on setting up local LLMs with Langroid, see here:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import json
import os
from typing import List

import chainlit as cl
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE, NO_ANSWER

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class LeasePeriod(BaseModel):
    start_date: str
    end_date: str


class LeaseFinancials(BaseModel):
    monthly_rent: str
    deposit: str


class Lease(BaseModel):
    """
    Various lease terms.
    Nested fields to make this more interesting/realistic
    """

    period: LeasePeriod
    financials: LeaseFinancials
    address: str


class LeaseMessage(ToolMessage):
    """Tool/function to use to present details about a commercial lease"""

    request: str = "lease_info"
    purpose: str = """
        Collect information about a Commercial Lease.
        """
    terms: Lease
    result: str = ""

    def handle(self) -> str:
        print(
            f"""
        DONE! Successfully extracted Lease Info:
        {self.terms}
        """
        )
        return "DONE " + json.dumps(self.terms.model_dump())

    @classmethod
    def format_instructions(cls, tool: bool = False) -> str:
        instr = super().format_instructions(tool)
        instr += """
        ------------------------------
        ASK ME QUESTIONS ONE BY ONE, to FILL IN THE FIELDS 
        of the `lease_info` function/tool.
        First ask me for the start date of the lease.
        DO NOT ASK ANYTHING ELSE UNTIL YOU RECEIVE MY ANSWER.
        """
        return instr

    @classmethod
    def examples(cls) -> List["LeaseMessage"]:
        return [
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-01-01", end_date="2021-12-31"),
                    financials=LeaseFinancials(monthly_rent="$1000", deposit="$1000"),
                    address="123 Main St, San Francisco, CA 94105",
                ),
                result="",
            ),
        ]


@cl.on_chat_start
async def main(
    debug: bool = False,
    model: str = "",  # or e.g. "ollama/nous-hermes2-mixtral",
    cache: bool = False,  # disables cache lookup; set to True to use cache
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=cache,
        )
    )
    llm_cfg = OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust based on model
        temperature=0,
        timeout=45,
    )
    doc_agent = DocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            n_neighbor_chunks=2,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=50,
                overlap=10,
                n_neighbor_ids=4,
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.vecdb.set_collection("docqa-chat-multi-extract", replace=True)
    print("[blue]Welcome to the real-estate info-extractor!")
    doc_agent.config.doc_paths = [
        "examples/docqa/lease.txt",
    ]
    doc_agent.ingest()
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        system_message="""You are an expert on Commercial Leases. 
        You will receive various questions about a Commercial 
        Lease contract, along with some excerpts from the Lease.
        Your job is to answer them concisely in at most 2 sentences.
        """,
    )

    lease_extractor_agent = ChatAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    lease_extractor_agent.enable_message(LeaseMessage)

    lease_task = Task(
        lease_extractor_agent,
        name="LeaseExtractor",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message=f"""
        You are an expert at understanding JSON function/tool specifications, and
        you are also very familiar with commercial lease terminology and concepts.
         
        See the `lease_info` function/tool below,  Your FINAL GOAL is to fill
        in the required fields in this `lease_info` function/tool,
        as shown in the example. This is ONLY an EXAMPLE,
        and YOU CANNOT MAKE UP VALUES FOR THESE FIELDS.
        
        To fill in these fields, you must ASK ME QUESTIONS about the lease,
        ONE BY ONE, and I will answer each question. 
        If I am unable to answer your question initially, try asking me 
        differently. If I am still unable to answer after 3 tries, fill in 
        {NO_ANSWER} for that field.
        When you have collected this info, present it to me using the 
        'lease_info' function/tool.
        DO NOT USE THIS Function/tool UNTIL YOU HAVE ASKED QUESTIONS 
        TO FILL IN ALL THE FIELDS.
        
        Think step by step. 
        Phrase each question simply as "What is ... ?",
        and do not explain yourself, or say any extraneous things. 
        Start by asking me for the start date of the lease.
        When you receive the answer, then ask for the next field, and so on.
        """,
    )

    validator_agent = ChatAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
            system_message=f"""
            You are obedient, understand instructions, and follow them faithfully,
            paying attention to the FORMAT specified,
            and you are also extremely CONCISE and SUCCINCT in your responses.
            
            Your task is to detect if the user's message is asking for ONE
            piece of information, or MULTIPLE pieces. Here is how you respond:
            
            IF the msg is only asking about ONE thing, OR if it is NOT EVEN a question:
                respond '{DONE}' and say nothing else.

            IF the msg is asking MORE THAN ONE thing,  respond like this:
            "Please only ask ONE question at a time. Try your question again.
            ONLY when you have ALL the answers, then present the info
            using the `lease_info` function/tool."
            """,
        )
    )
    validator_task = Task(
        validator_agent,
        name="Validator",
        single_round=True,
        interactive=False,
    )

    lease_task.add_sub_task([validator_task, doc_task])
    lr.ChainlitTaskCallbacks(lease_task)

    # DocChatAgent does not have an async llm_response method,
    # so we must use task.run() instead of task.run_async(),
    # but fortunately we can wrap it in a cl.make_async() call
    await cl.make_async(lease_task.run)()
</file>

<file path="examples/chainlit/multi-extract.py">
"""
TODO: examples does not work due to breaking changes in Chainlit

Two-agent chat with Retrieval-augmented LLM + function-call/tool.
ExtractorAgent (has no access to docs) is tasked with extracting structured
information from a commercial lease document, and must present the terms in
a specific nested JSON format.
This agent generates questions corresponding to each field in the JSON format,
and the RAG-enabled DocAgent (has access to the lease) answers the  questions.

This is a Chainlit version of examples/docqa/chat_multi_extract.py.

Example:
chainlit run examples/chainlit/multi-extract.py

This uses GPT4-turbo by default, but works very well with the `dolphin-mixtral`
local LLM, which you can specify in the llm_config below
using `chat_model = "ollama/dolphin-mixtral:latest"`,
provided you've already spun it up with ollama:
```
ollama run dolphin-mixtral
```

See here for more on setting up LLMs to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

The challenging parts in this script are agent-to-agent delegation, and the extractor
agent planning out a sequence of questions to ask the doc agent, and finally presenting
the collected information in a structured format to the user using a Tool/Function-call.
The `dolphin-mixtral` model seems to handle this pretty well, however weaker models
may not be able to handle this.

"""

import json
import os
from typing import List

import chainlit as cl
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class LeasePeriod(BaseModel):
    start_date: str
    end_date: str


class LeaseFinancials(BaseModel):
    monthly_rent: str
    deposit: str


class Lease(BaseModel):
    """
    Various lease terms.
    Nested fields to make this more interesting/realistic
    """

    period: LeasePeriod
    financials: LeaseFinancials
    address: str


class LeaseMessage(ToolMessage):
    """Tool/function to use to present details about a commercial lease"""

    request: str = "lease_info"
    purpose: str = """
        Collect information about a Commercial Lease.
        """
    terms: Lease
    result: str = ""

    @classmethod
    def examples(cls) -> List["LeaseMessage"]:
        return [
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-01-01", end_date="2021-12-31"),
                    financials=LeaseFinancials(monthly_rent="$1000", deposit="$1000"),
                    address="123 Main St, San Francisco, CA 94105",
                ),
                result="",
            ),
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-04-01", end_date="2022-04-28"),
                    financials=LeaseFinancials(monthly_rent="$2000", deposit="$2000"),
                    address="456 Main St, San Francisco, CA 94111",
                ),
                result="",
            ),
        ]


class LeaseExtractorAgent(ChatAgent):
    def __init__(self, config: ChatAgentConfig):
        super().__init__(config)

    def lease_info(self, message: LeaseMessage) -> str:
        print(
            f"""
        DONE! Successfully extracted Lease Info:
        {message.terms}
        """
        )
        return "DONE \n" + json.dumps(message.terms.model_dump(), indent=4)


@cl.on_chat_start
async def main(
    debug: bool = False,
    model: str = "",  # or "ollama/dolphin-mixtral:latest"
    nocache: bool = False,
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    llm_cfg = OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust based on model
        temperature=0,
        timeout=45,
    )
    doc_agent = DocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=300,
                overlap=50,
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.vecdb.set_collection("docqa-chat-multi-extract", replace=True)
    print("[blue]Welcome to the real-estate info-extractor!")
    doc_agent.config.doc_paths = [
        "examples/docqa/lease.txt",
    ]
    doc_agent.ingest()
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        system_message="""You are an expert on Commercial Leases. 
        You will receive various questions about a Commercial 
        Lease contract, along with some excerpts from the Lease.
        Your job is to answer them concisely in at most 2 sentences.
        """,
    )

    lease_extractor_agent = LeaseExtractorAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    lease_extractor_agent.enable_message(LeaseMessage)

    lease_task = Task(
        lease_extractor_agent,
        name="LeaseExtractor",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message=f"""
        You have to collect some SPECIFIC STRUCTURED information 
        about a Commercial Lease, as specified in the `lease_info` function/tool. 
        But you do not have access to the lease itself. 
        You can ask me questions about the lease, ONE AT A TIME, I will answer each 
        question. You only need to collect info to fill the fields in the 
        `field_info` function/tool. 
        If I am unable to answer your question initially, try asking me 
        differently. If I am still unable to answer after 3 tries, fill in 
        {NO_ANSWER} for that field.
        When you have collected this info, present it to me using the 
        'lease_info' function/tool.
        DO NOT USE THIS Function/tool UNTIL YOU HAVE ASKED QUESTIONS 
        TO FILL IN ALL THE FIELDS.
        
        Start by asking me for the start date of the lease.
        """,
    )
    lease_task.add_sub_task(doc_task)
    # The below line is essentially the ONLY change to make
    # to the original script on which this is based.
    lr.ChainlitTaskCallbacks(lease_task)
    # DocChatAgent does not have an async llm_response method,
    # so we must use task.run() instead of task.run_async(),
    # but fortunately we can wrap it in a cl.make_async() call
    await cl.make_async(lease_task.run)()
</file>

<file path="examples/chainlit/README.md">
# Running the chainlit apps

In your Python virtual env, ensure you have 
installed `langroid` with the `chainlit` extra using, e.g.

```bash
pip install langroid[chainlit]
```

Or if you already have `langroid` installed, you can install the `chainlit` extra using:

```bash
pip install chainlit
```

To check that `chainlit` is installed, run:

```bash
chainlit hello
```

and you should see the `hello app` open in your browser.


## General usage
See [chainlit docs](https://docs.chainlit.io/get-started/overview) to learn the basics.

Generally speaking to use Langroid `ChatAgents` or `Tasks` with 
`chainlit`, you simply need to wrap your `ChatAgent` or `Task` in the appropriate 
"callback injection" class, e.g. either
```
import langroid as lr
agent = lr.ChatAgent(...)
lr.ChainlitAgentCallbacks(agent) 
```
or 
```
task = lr.Task(...)
lr.ChainlitTaskCallbacks(task) 
```
The `ChainlitTaskCallbacks` class recursively injects callbacks into 
`ChatAgents` belonging to the task, and any sub-tasks.
The callback classes are defined 
[here](https://github.com/langroid/langroid/blob/main/langroid/agent/callbacks/chainlit.py).

You also need to write an `on_chat_start` function and possibly an `on_message`
function to start off the app. See the examples to learn more.

## Configuration

⚠️ It is very important that you download the `.chainlit` directory from the `langroid` repo
(or the `langroid-examples` repo) and place it *in the directory from
which you run the `chainlit` command*. E.g. if you run the `chainlit` command from the
root of the repo, then the `.chainlit` directory should be placed there.
This directory contains various customizations, but most importantly, it contains the
file `translations/en-US.json`, where the default placeholder text in the chat box is defined
(as described below as well). If you've correctly placed this directory, this default text should say
something like 
```
Ask, respond, give feedback, or just 'c' for continue...
```

You can configure some aspects of the chainlit app via these files,
which are included in this repo at the root level (see
the Chainlit [customization docs](https://docs.chainlit.io/customisation/overview) for more details):
- `.chainlit/config.toml` to customize project, features, UI (see [here](https://docs.chainlit.io/backend/config/overview))
- `.chainlit/translations/en-US.json` for various ["translations"](https://docs.chainlit.io/customisation/translation) and language-specific
   customizations. In particular, the default text in the input box is customized here.
- `chainlit.md`, which contains the initial "Readme" content
- [Logo, favicons](https://docs.chainlit.io/customisation/custom-logo-and-favicon) should be placed in a directory
  named `public` adjacent to the apps. 

Depending on how you organize your apps, you may need to run the `chainlit` command 
from the directory where the above customization files/dirs are placed.
</file>

<file path="examples/chainlit/simplest.py">
"""
Absolute bare-bones way to set up a simple chatbot using all default settings,
using a Langroid Task + callbacks.

After setting up the virtual env as in README,
and you have your OpenAI API Key in the .env file, run like this:

chainlit run examples/chainlit/simplest.py
"""

import chainlit as cl

import langroid as lr
import langroid.language_models as lm


@cl.on_message
async def on_message(message: cl.Message):
    lm_config = lm.OpenAIGPTConfig()
    agent = lr.ChatAgent(lr.ChatAgentConfig(llm=lm_config))
    task = lr.Task(agent, interactive=True)

    lr.ChainlitTaskCallbacks(task)
    await task.run_async(message.content)
</file>

<file path="examples/chainlit/test-step-nesting.py">
"""
Test whether the current chainlit version shows nested steps as expected.
Note that this does NOT show what you'd expect, due to breaking changes in Chainlit.

Two things to look for:
(1) are all types of steps shown, or only type = "tool"?
(2) when step B has parent_id pointing to Step A, we want to see Step B shown:
    - nested under Step A
    - shown in a chronologically correct order, i.e. if Step A says "hello",
        then calls Step B, then step B should be shown AFTER the "hello" message from A.

(1) is fine in chainlit 1.1.202, i.e. all steps are shown whether tools or not
    but in 1.1.300, only type = "tool" steps are shown.
    For example if the `type` params are other than "tool" in the example below,
    the steps will not show up in the chat.
(2) is broken in 1.1.202 -- the sub-step is correctly nested BUT always shows up
    at the TOP, and can look very unintuitive, as this example shows.
"""

import chainlit as cl


@cl.on_chat_start
async def on_chat_start():
    a_step = cl.Step(name="A", type="tool")
    a_step.output = "asking B"
    await a_step.send()

    b_step = cl.Step(
        name="B",
        parent_id=a_step.id,
        type="tool",
    )
    b_step.output = "asking C"
    await b_step.send()

    c_step = cl.Step(
        name="C",
        parent_id=b_step.id,
        type="tool",
    )
    c_step.output = "C answered!"
    await c_step.send()
</file>

<file path="examples/data-qa/sql-chat/__init__.py">

</file>

<file path="examples/data-qa/sql-chat/sql_chat.py">
"""
Example showing how to chat with a SQL database.

Note if you are using this with a postgres db, you will need to:

(a) Install PostgreSQL dev libraries for your platform, e.g.
    - `sudo apt-get install libpq-dev` on Ubuntu,
    - `brew install postgresql` on Mac, etc.
(b) langroid with the postgres extra, e.g. `pip install langroid[postgres]`
    or `poetry add langroid[postgres]` or `poetry install -E postgres`
    or `uv pip install langroid[postgres]` or `uv add langroid[postgres]`.
    If this gives you an error, try `pip install psycopg2-binary` in your virtualenv.
"""

import json
import os
from typing import Any, Dict

import typer
from rich import print
from rich.prompt import Prompt

from langroid.exceptions import LangroidImportError

try:
    from sqlalchemy import create_engine, inspect
    from sqlalchemy.engine import Engine
except ImportError as e:
    raise LangroidImportError(extra="sql", error=str(e))

from prettytable import PrettyTable

try:
    from .utils import fix_uri, get_database_uri
except ImportError:
    from utils import fix_uri, get_database_uri
import logging

from langroid.agent.special.sql.sql_chat_agent import (
    SQLChatAgent,
    SQLChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import SEND_TO

logger = logging.getLogger(__name__)


app = typer.Typer()


def create_descriptions_file(filepath: str, engine: Engine) -> None:
    """
    Create an empty descriptions JSON file for SQLAlchemy tables.

    This function inspects the database, generates a template for table and
    column descriptions, and writes that template to a new JSON file.

    Args:
        filepath: The path to the file where the descriptions should be written.
        engine: The SQLAlchemy Engine connected to the database to describe.

    Raises:
        FileExistsError: If the file at `filepath` already exists.

    Returns:
        None
    """
    if os.path.exists(filepath):
        raise FileExistsError(f"File {filepath} already exists.")

    inspector = inspect(engine)
    descriptions: Dict[str, Dict[str, Any]] = {}

    for table_name in inspector.get_table_names():
        descriptions[table_name] = {
            "description": "",
            "columns": {col["name"]: "" for col in inspector.get_columns(table_name)},
        }

    with open(filepath, "w") as f:
        json.dump(descriptions, f, indent=4)


def load_context_descriptions(engine: Engine) -> dict:
    """
    Ask the user for a path to a JSON file and load context descriptions from it.

    Returns:
        dict: The context descriptions, or an empty dictionary if the user decides to skip this step.
    """

    while True:
        filepath = Prompt.ask(
            "[blue]Enter the path to your context descriptions file. \n"
            "('n' to create a NEW file, 's' to SKIP, or Hit enter to use DEFAULT) ",
            default="examples/data-qa/sql-chat/demo.json",
        )

        if filepath.strip() == "s":
            return {}

        if filepath.strip() == "n":
            filepath = Prompt.ask(
                "[blue]To create a new context description file, enter the path",
                default="examples/data-qa/sql-chat/description.json",
            )
            print(f"[blue]Creating new context description file at {filepath}...")
            create_descriptions_file(filepath, engine)
            print(
                f"[blue] Please fill in the descriptions in {filepath}, "
                f"then try again."
            )

        # Try to load the file
        if not os.path.exists(filepath):
            print(f"[red]The file '{filepath}' does not exist. Please try again.")
            continue

        try:
            with open(filepath, "r") as file:
                return json.load(file)
        except json.JSONDecodeError:
            print(
                f"[red]The file '{filepath}' is not a valid JSON file. Please try again."
            )


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    tools: bool = typer.Option(
        False, "--tools", "-t", help="use langroid tools instead of function-calling"
    ),
    schema_tools: bool = typer.Option(
        False, "--schema_tools", "-st", help="use schema tools"
    ),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
            cache_type="redis",
        )
    )
    print("[blue]Welcome to the SQL database chatbot!\n")
    database_uri = Prompt.ask(
        """
        [blue]Enter the URI for your SQL database 
        (type 'i' for interactive, or hit enter for default)
        """,
        default="sqlite:///examples/data-qa/sql-chat/demo.db",
    )

    if database_uri == "i":
        database_uri = get_database_uri()

    database_uri = fix_uri(database_uri)
    logger.warning(f"Using database URI: {database_uri}")

    # Create engine and inspector
    engine = create_engine(database_uri)
    inspector = inspect(engine)

    context_descriptions = load_context_descriptions(engine)

    # Get table names
    table_names = inspector.get_table_names()

    for table_name in table_names:
        print(f"[blue]Table: {table_name}")

        # Create a new table for the columns
        table = PrettyTable()
        table.field_names = ["Column Name", "Type"]

        # Get the columns for the table
        columns = inspector.get_columns(table_name)
        for column in columns:
            table.add_row([column["name"], column["type"]])

        print(table)

    agent_config = SQLChatAgentConfig(
        name="sql",
        database_uri=database_uri,
        use_tools=tools,
        use_functions_api=not tools,
        show_stats=False,
        chat_mode=True,
        use_helper=True,
        context_descriptions=context_descriptions,  # Add context descriptions to the config
        use_schema_tools=schema_tools,
        addressing_prefix=SEND_TO,
        llm=OpenAIGPTConfig(
            chat_model=OpenAIChatModel.GPT4o,
        ),
    )
    agent = SQLChatAgent(agent_config)
    # Set interactive = False, but we user gets chance to respond
    # when explicitly addressed by LLM
    task = Task(agent, interactive=False)
    task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/data-qa/sql-chat/utils.py">
import logging
import urllib.parse

from rich import print
from rich.prompt import Prompt

from langroid.parsing.utils import closest_string

logger = logging.getLogger(__name__)


DEFAULT_PORTS = dict(
    postgresql=5432,
    mysql=3306,
    mariadb=3306,
    mssql=1433,
    oracle=1521,
    mongodb=27017,
    redis=6379,
)


def fix_uri(uri: str) -> str:
    """Fixes a URI by percent-encoding the username and password."""

    if "%" in uri:
        return uri  # already %-encoded, so don't do anything
    # Split by '://'
    scheme_part, rest_of_uri = uri.split("://", 1)

    # Get the final '@' (assuming only the last '@' is the separator for user info)
    last_at_index = rest_of_uri.rfind("@")
    userinfo_part = rest_of_uri[:last_at_index]
    rest_of_uri_after_at = rest_of_uri[last_at_index + 1 :]

    if ":" not in userinfo_part:
        return uri
    # Split userinfo by ':' to get username and password
    username, password = userinfo_part.split(":", 1)

    # Percent-encode the username and password
    username = urllib.parse.quote(username)
    password = urllib.parse.quote(password)

    # Construct the fixed URI
    fixed_uri = f"{scheme_part}://{username}:{password}@{rest_of_uri_after_at}"

    return fixed_uri


def _create_database_uri(
    scheme: str,
    username: str,
    password: str,
    hostname: str,
    port: int,
    databasename: str,
) -> str:
    """Generates a database URI based on provided parameters."""
    username = urllib.parse.quote_plus(username)
    password = urllib.parse.quote_plus(password)
    port_str = f":{port}" if port else ""
    return f"{scheme}://{username}:{password}@{hostname}{port_str}/{databasename}"


def get_database_uri() -> str:
    """Main function to gather input and print the database URI."""
    scheme_input = Prompt.ask("Enter the database type (e.g., postgresql, mysql)")
    scheme = closest_string(scheme_input, list(DEFAULT_PORTS.keys()))

    # Handle if no close match is found.
    if scheme == "No match found":
        print(f"No close match found for '{scheme_input}'. Please verify your input.")
        return

    username = Prompt.ask("Enter the database username")
    password = Prompt.ask("Enter the database password", password=True)
    hostname = Prompt.ask("Enter the database hostname")

    # Inform user of default port, and let them choose to override or leave blank
    default_port = DEFAULT_PORTS.get(scheme, "")
    port_msg = (
        f"Enter the database port "
        f"(hit enter to use default: {default_port} or specify another value)"
    )

    port = Prompt.ask(port_msg, default=default_port)
    if not port:  # If user pressed enter without entering anything
        port = default_port
    port = int(port)

    databasename = Prompt.ask("Enter the database name")

    uri = _create_database_uri(scheme, username, password, hostname, port, databasename)
    print(f"Your {scheme.upper()} URI is:\n{uri}")
    return uri
</file>

<file path="examples/data-qa/__init__.py">

</file>

<file path="examples/data-qa/table_chat.py">
"""
Example showing how to chat with a tabular dataset:
csv, tsv, or any other pandas-readable.

Run like this

python3 examples/data-qa/table_chat.py

Optional args:
* -d or --debug to enable debug mode
* -ns or --nostream to disable streaming
* -nc or --nocache to disable caching
* -m or --model to specify a model name

To run with a local model via ollama, do this:
```
ollama run dolphin-mixtral # best model for this script

python3 examples/data-qa/table_chat.py -m ollama/dolphin-mixtral:latest
```

For more info on running Langroid with local LLM, see here:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import typer
from rich import print
from rich.prompt import Prompt

from langroid.agent.special.table_chat_agent import TableChatAgent, TableChatAgentConfig
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )

    print("[blue]Welcome to the tabular-data chatbot!\n")
    path = Prompt.ask(
        "[blue]Enter a local path or URL to a tabular dataset (hit enter to use default)\n",
        default="https://raw.githubusercontent.com/fivethirtyeight/data/master/airline-safety/airline-safety.csv",
    )

    agent = TableChatAgent(
        config=TableChatAgentConfig(
            data=path,
            llm=OpenAIGPTConfig(
                chat_model=model or OpenAIChatModel.GPT4o,
                chat_context_length=16_000,  # adjust based on model
                timeout=45,
                temperature=0.2,
            ),
        )
    )
    task = Task(agent, interactive=True)
    task.run("Can you help me with some questions about a tabular dataset?")


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/streamlit-app/app.py">
import os

import streamlit as st
from utils import agent, configure

import langroid.language_models as lm
from langroid.utils.configuration import settings

settings.cache_type = "fakeredis"
if st.session_state.get("specified_file") is None:
    st.session_state["specified_file"] = ""
if st.session_state.get("file_path") is None:
    st.session_state["file_path"] = ""
if st.session_state.get("rag_agent") is None:
    st.session_state["rag_agent"] = None
if st.session_state.get("chat_model") is None:
    st.session_state["chat_model"] = None

default_chat_model = lm.OpenAIChatModel.GPT4o.value
chat_model = st.sidebar.text_input(
    f"""
Chat model, e.g. `litellm/ollama/mistral:7b-instruct-v0.2-q4_K_M`,
or leave empty to default to {default_chat_model}
"""
)
actual_chat_model = chat_model or default_chat_model
st.session_state["chat_model"] = actual_chat_model
st.sidebar.info(f"Using chat model: {str(actual_chat_model)}")
st.header("DocChatAgent by Langroid", divider="rainbow")

uploaded_file = st.file_uploader("Choose a txt file")
TEMP_DIR = "tempdir"
if uploaded_file is not None:
    if uploaded_file.name != st.session_state["specified_file"]:
        temp_dir = os.makedirs(TEMP_DIR, exist_ok=True)
        temp_path = os.path.join(TEMP_DIR, uploaded_file.name)
        with open(temp_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.session_state["specified_file"] = uploaded_file.name
        st.session_state["file_path"] = temp_path
    else:
        temp_path = st.session_state["file_path"]

temp_path = st.session_state["file_path"]
cfg = configure(temp_path, actual_chat_model)

prompt = st.chat_input("Talk with Document")
if prompt:
    st.write(f"{prompt}")

    # chat using docchatagent
    answer = agent(cfg, prompt)
    st.write(f"{answer}")
</file>

<file path="examples/docqa/streamlit-app/README.md">
# Basic example: chat with a document using Langroid with local LLM or OpenAI LLM

Bare-bones example of an app that combines:
- Langroid `DocChatAgent` for RAG
- StreamLit for webapp/UI
to let you ask questions about the contents of a file (pdf, txt, docx, md, html).

## Instructions
Run this from the root of the `langroid-examples` repo. Assuming you already have a virtual env in 
which you have installed `langroid`, the only additional requirement is to run:

``` 
pip install streamlit
```
Then run the application like this:
```
streamlit run examples/docqa/streamlit-app/app.py
```
In the sidebar you can specify a local LLM, or leave it blank to use the OpenAI 
GPT4-Turbo model. 


## Limitations

- Streaming does not currently work
- Conversation is not accumulated
- Source, Extract evidence-citation is only displayed in terminal/console, to reduce clutter in the UI.

## Credits
Code adapted from Prashant Kumar's example in [`lancedb/vectordb-recipies`](https://github.com/lancedb/vectordb-recipes)
</file>

<file path="examples/docqa/streamlit-app/requirements.txt">
langroid
streamlit
</file>

<file path="examples/docqa/streamlit-app/utils.py">
import os

import streamlit as st

from langroid.agent.special import DocChatAgent, DocChatAgentConfig
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.parsing.parser import ParsingConfig
from langroid.vector_store.qdrantdb import QdrantDBConfig

OPENAI_KEY = os.environ["OPENAI_API_KEY"]


@st.cache_data
def configure(filename: str, chat_model: str = "") -> DocChatAgentConfig:
    llm_cfg = OpenAIGPTConfig(
        chat_model=chat_model,
    )

    oai_embed_config = OpenAIEmbeddingsConfig(
        model_type="openai",
        model_name="text-embedding-3-small",
        dims=1536,
    )

    # Configuring DocChatAgent
    cfg = DocChatAgentConfig(
        n_similar_chunks=4,
        n_relevant_chunks=4,
        parsing=ParsingConfig(
            chunk_size=100,
            overlap=20,
        ),
        show_stats=False,
        cross_encoder_reranking_model="",
        llm=llm_cfg,
        vecdb=QdrantDBConfig(
            embedding=oai_embed_config,
            collection_name="lease",
            replace_collection=True,
            cloud=False,
        ),
        doc_paths=[filename],
    )

    return cfg


def agent(cfg, prompt):
    # Creating DocChatAgent
    rag_agent = st.session_state["rag_agent"]
    if (
        rag_agent is None
        or st.session_state["chat_model"] != cfg.llm.chat_model
        or st.session_state["file_path"] != cfg.doc_paths[0]
    ):
        rag_agent = DocChatAgent(cfg)
        st.session_state["rag_agent"] = rag_agent

    response = rag_agent.llm_response(prompt)
    return response.content
</file>

<file path="examples/docqa/__init__.py">

</file>

<file path="examples/docqa/books.txt">
Book Title: Crime and Redemption by Filidor Dostoyevski, released in 1877, offers a
riveting exploration of guilt, morality, and the possibility of spiritual rebirth.
Set against the bleak backdrop of 19th century Russia, it follows the tormented journey
of Rodion Romanovich Raskolnikov, a young man driven to murder and subsequently
haunted by his actions. Through Raskolnikov's story, Dostoyevski delves deep into the
human psyche, presenting a timeless narrative of human imperfection and the
redemptive power.

Book Title: The Siblings Karamazoff by Fyodar Dostoyevskiy, published in 1881,
weaves a complex narrative around the ethical battles and spiritual dilemmas
faced by the Karamazoff family. Set in the heart of Russia, it explores themes of faith,
doubt, and the nature of free will through the intersecting lives of three brothers,
each embodying different facets of humanity. Dostoyevskiy masterfully crafts a tale of
familial bonds, existential questioning, and the search for truth in a morally ambiguous
world.
</file>

<file path="examples/docqa/chat_multi_extract.py">
"""
Two-agent chat with Retrieval-augmented LLM + function-call/tool.
ExtractorAgent (has no access to docs) is tasked with extracting structured
information from a commercial lease document, and must present the terms in
a specific nested JSON format.
This agent generates questions corresponding to each field in the JSON format,
and the RAG-enabled DocAgent (has access to the lease) answers the  questions.


Example:
python3 examples/docqa/chat_multi_extract.py

This uses a GPT4 model by default, but works very well with the `dolphin-mixtral`
local LLM, which you can specify via the -m arg:

```
ollama run dolphin-mixtral

python3 examples/docqa/chat_multi_extract.py -m ollama/dolphin-mixtral:latest
```

The challenging parts in this script are agent-to-agent delegation, and the extractor
agent planning out a sequence of questions to ask the doc agent, and finally presenting
the collected information in a structured format to the user using a Tool/Function-call.
The `dolphin-mixtral` model seems to handle this pretty well, however weaker models
may not be able to handle this.

For weaker LLMs, the script examples/docqa/chat-multi-extract-local.py performs a similar task
but uses a workflow where agents do not delegate to each other,
and uses more agents to break down tasks into smaller parts.

"""

import json
import os
from typing import List

import typer
from rich import print

import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class LeasePeriod(BaseModel):
    start_date: str
    end_date: str


class LeaseFinancials(BaseModel):
    monthly_rent: str
    deposit: str


class Lease(BaseModel):
    """
    Various lease terms.
    Nested fields to make this more interesting/realistic
    """

    period: LeasePeriod
    financials: LeaseFinancials
    address: str


class LeaseMessage(ToolMessage):
    """Tool/function to use to present details about a commercial lease"""

    request: str = "lease_info"
    purpose: str = """
        Collect information about a Commercial Lease.
        """
    terms: Lease
    result: str = ""

    @classmethod
    def examples(cls) -> List["LeaseMessage"]:
        return [
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-01-01", end_date="2021-12-31"),
                    financials=LeaseFinancials(monthly_rent="$1000", deposit="$1000"),
                    address="123 Main St, San Francisco, CA 94105",
                ),
                result="",
            ),
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-04-01", end_date="2022-04-28"),
                    financials=LeaseFinancials(monthly_rent="$2000", deposit="$2000"),
                    address="456 Main St, San Francisco, CA 94111",
                ),
                result="",
            ),
        ]


class LeaseExtractorAgent(ChatAgent):
    def __init__(self, config: ChatAgentConfig):
        super().__init__(config)

    def lease_info(self, message: LeaseMessage) -> str:
        print(
            f"""
        DONE! Successfully extracted Lease Info:
        {message.terms}
        """
        )
        return "DONE " + json.dumps(message.terms.model_dump())


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    llm_cfg = OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust based on model
        temperature=0,
        timeout=45,
    )
    doc_agent = DocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=300,
                overlap=50,
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.vecdb.set_collection("docqa-chat-multi-extract", replace=True)
    print("[blue]Welcome to the real-estate info-extractor!")
    doc_agent.config.doc_paths = [
        "examples/docqa/lease.txt",
    ]
    doc_agent.ingest()
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        system_message="""You are an expert on Commercial Leases. 
        You will receive various questions about a Commercial 
        Lease contract, along with some excerpts from the Lease.
        Your job is to answer them concisely in at most 2 sentences.
        """,
    )

    lease_extractor_agent = LeaseExtractorAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    lease_extractor_agent.enable_message(LeaseMessage)

    lease_task = Task(
        lease_extractor_agent,
        name="LeaseExtractorAgent",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message=f"""
        You have to collect some SPECIFIC STRUCTURED information 
        about a Commercial Lease, as specified in the `lease_info` function/tool. 
        But you do not have access to the lease itself. 
        You can ask me questions about the lease, ONE AT A TIME, I will answer each 
        question. You only need to collect info to fill the fields in the 
        `lease_info` function/tool. 
        If I am unable to answer your question initially, try asking me 
        differently. If I am still unable to answer after 3 tries, fill in 
        {NO_ANSWER} for that field.
        When you have collected this info, present it to me using the 
        'lease_info' function/tool.
        DO NOT USE THIS Function/tool UNTIL YOU HAVE ASKED QUESTIONS 
        TO FILL IN ALL THE FIELDS.
        
        Start by asking me for the start date of the lease.
        """,
    )
    lease_task.add_sub_task(doc_task)
    lease_task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/chat_search.py">
"""
This is a single-agent question-answering system that has access to a Web-Search
Tool when needed,
and in case a web search is used, ingests scraped link contents into a vector-db,
and uses Retrieval Augmentation to answer the question.

Run like this:

    python3 examples/docqa/chat-search.py -m groq/llama-3.1-70b-versatile

The -m arg is optional, defaults to GPT4o

Optional args:
    -nc : turn off caching (i.e. don't retrieve cached LLM responses)
    -d: debug mode, to show all intermediate results
    -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", "exa", "crawl4ai"

See here for guide to using local LLMs with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import logging
import re
from typing import Any, List, Optional

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatDocument
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import ForwardTool
from langroid.parsing.url_loader import (
    ExaCrawlerConfig,
    FirecrawlConfig,
    TrafilaturaConfig,
    Crawl4aiConfig,
)
from langroid.parsing.web_search import exa_search
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

logger = logging.getLogger(__name__)


class RelevantExtractsTool(ToolMessage):
    request: str = "relevant_extracts"
    purpose: str = "Get docs/extracts relevant to the <query>"
    query: str

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(query="when was the Mistral LLM released?"),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """


class RelevantSearchExtractsTool(ToolMessage):
    request: str = "relevant_search_extracts"
    purpose: str = "Get docs/extracts relevant to the <query> from a web search"
    query: str
    num_results: int = 3

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(
                query="when was the Mistral LLM released?",
                num_results=3,
            ),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """


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 == "trafilatura" or crawler is None:
            self.config.crawler_config = TrafilaturaConfig()
        elif crawler == "firecrawl":
            self.config.crawler_config = FirecrawlConfig()
        elif crawler == "exa":
            self.config.crawler_config = ExaCrawlerConfig()
        elif crawler == "crawl4ai":
            self.config.crawler_config = Crawl4aiConfig()
        else:
            raise ValueError(
                f"Unsupported crawler {crawler}. Options are: 'trafilatura', 'firecrawl', 'exa', 'crawl4ai'"
            )

    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> ChatDocument | None:
        # override llm_response of DocChatAgent to allow use of the tools.
        return ChatAgent.llm_response(self, message)

    def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            return ForwardTool(agent="user")

    def relevant_extracts(self, msg: RelevantExtractsTool) -> str:
        """Get docs/extracts relevant to the query, from vecdb"""
        self.tried_vecdb = True
        query = msg.query
        _, extracts = self.get_relevant_extracts(query)
        if len(extracts) == 0:
            return """
            No extracts found! You can try doing a web search with the
            `relevant_search_extracts` tool/function-call.
            """
        return "\n".join(str(e) for e in extracts)

    def relevant_search_extracts(self, msg: RelevantSearchExtractsTool) -> str:
        """Get docs/extracts relevant to the query, from a web search"""
        if not self.tried_vecdb and len(self.original_docs) > 0:
            return "Please try the `relevant_extracts` tool, before using this tool"
        self.tried_vecdb = False
        query = msg.query
        num_results = msg.num_results
        logger.warning("Trying exa search...")
        results = exa_search(query, num_results)
        links = [r.link for r in results]
        logger.warning(f"Found {len(links)} links, ingesting into vecdb...")
        self.config.doc_paths = links
        self.ingest()
        logger.warning(f"Ingested {len(links)} links into vecdb")
        _, extracts = self.get_relevant_extracts(query)
        return "\n".join(str(e) for e in extracts)


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = "",
    fn_api: bool = True,
    crawler: Optional[str] = None,
) -> 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 (Optional[str]): Specify the crawler to use for web search.
                                Options are: trafilatura (default), firecrawl, exa, crawl4ai.
    """

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    print(
        """
        [blue]Welcome to the Internet Search chatbot!
        I will try to answer your questions, relying on (full content of links from) 
        Duckduckgo (DDG) Search when needed.
        
        Enter x or q to quit, or ? for evidence
        """
    )

    system_msg = Prompt.ask(
        """
    [blue] Tell me who I am (give me a role) by completing this sentence: 
    You are...
    [or hit enter for default]
    [blue] Human
    """,
        default="a helpful assistant.",
    )
    system_msg = re.sub("you are", "", system_msg, flags=re.IGNORECASE)

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=16_000,  # adjust based on model
    )

    config = DocChatAgentConfig(
        use_functions_api=fn_api,
        use_tools=not fn_api,
        llm=llm_config,
        system_message=f"""
        {system_msg} You will try your best to answer my questions,
        in this order of preference:
        1. If you can answer from your own knowledge, simply return the answer
        2. Otherwise, ask me for some relevant text, and I will send you. Use the 
            `relevant_extracts` tool/function-call for this purpose. Once you receive 
            the text, you can use it to answer my question. 
            If I say {NO_ANSWER}, it means I found no relevant docs, and you can try 
            the next step, using a web search.
        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.
        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`.
        
        Be very concise in your responses, use no more than 1-2 sentences.
        When you answer based on provided documents, be sure to show me 
        the SOURCE(s) and EXTRACT(s), for example:
        
        SOURCE: https://www.wikihow.com/Be-a-Good-Assistant-Manager
        EXTRACT: Be a Good Assistant ... requires good leadership skills.
        
        For the EXTRACT, ONLY show up to first 3 words, and last 3 words.
        """,
    )

    agent = SearchDocChatAgent(config, crawler=crawler)
    agent.enable_message(
        [
            RelevantExtractsTool,
            RelevantSearchExtractsTool,
        ]
    )
    collection_name = Prompt.ask(
        "Name a collection to use",
        default="docqa-chat-search",
    )
    replace = (
        Prompt.ask(
            "Would you like to replace (i.e. erase) this collection?",
            choices=["y", "n"],
            default="n",
        )
        == "y"
    )

    print(f"[red]Using {collection_name}")

    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 and crawling?"
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/chat-local.py">
"""
Single agent to use to chat with an LLM using  Retrieval-Augmented Generation (RAG).
Similar to chat.py but allows specifying a local LLM.

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

NOTES:
(1) The app works best with GPT4/Turbo, but results may be mixed with local LLMs.
You may have to tweak the system_message, use_message, and summarize_prompt
as indicated in comments below, to get good results.
(2) The default vector-db in DocChatAgent is QdrantDB, but you can switch to the
other supported vector-dbs, e.g. lancedb or chroma.

"""

import os
import re

import typer
from rich import print
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=32_000,  # adjust based on model
        timeout=90,
    )

    config = DocChatAgentConfig(
        n_query_rephrases=0,
        hypothetical_answer=False,
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=0,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        llm=llm_config,
        # relevance_extractor_config=None,
        # system_message="...override default DocChatAgent system msg here",
        # user_message="...override default DocChatAgent user msg here",
        # summarize_prompt="...override default DocChatAgent summarize prompt here",
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=300,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    agent = DocChatAgent(config)
    print("[blue]Welcome to the document chatbot!")
    agent.user_docs_ingest_dialog()
    print("[cyan]Enter x or q to quit, or ? for evidence")

    system_msg = Prompt.ask(
        """
    [blue] Tell me who I am; complete this sentence: You are...
    [or hit enter for default] 
    [blue] Human
    """,
        default="a helpful assistant.",
    )
    system_msg = re.sub("you are", "", system_msg, flags=re.IGNORECASE)
    task = Task(
        agent,
        system_message="You are " + system_msg,
    )
    task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/chat-multi-extract-3.py">
"""
Variant of chat_multi_extract.py more suited to local LLM, using 3 Agents
(instead of 2 agents):

- LeaseExtractorAgent: is tasked with extracting structured information from a
    commercial lease document, and must present the terms in a specific nested JSON
    format. This agent generates questions corresponding to each field in the JSON
    format.
- Validator: This agent detects if LeaseExtractorAgent's message is asking for ONE
    piece of information, or MULTIPLE pieces. If the message is only asking about ONE
    thing, OR if it is NOT EVEN a question, it responds with "DONE" and says nothing.
    If the message is asking MORE THAN ONE thing, it responds with a message asking to
    only ask ONE question at a time.
    [Why restrict to one question at a time? Because the DocAgent is more likely to
      understand and answer a single question at a time]

- DocAgent: This agent answers the questions generated by LeaseExtractorAgent,
    based on the lease document it has access to via vecdb, using RAG.

Run like this:

```
python3 examples/docqa/chat-multi-extract-3.py -m ollama/nous-hermes2-mixtral
```

If you omit the -m arg, it will use the default GPT4-turbo model.

For more on setting up local LLMs with Langroid, see here:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import json
import os
from typing import List

import typer
from rich import print

import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE, NO_ANSWER

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class LeasePeriod(BaseModel):
    start_date: str
    end_date: str


class LeaseFinancials(BaseModel):
    monthly_rent: str
    deposit: str


class Lease(BaseModel):
    """
    Various lease terms.
    Nested fields to make this more interesting/realistic
    """

    period: LeasePeriod
    financials: LeaseFinancials
    address: str


class LeaseMessage(ToolMessage):
    """Tool/function to use to present details about a commercial lease"""

    request: str = "lease_info"
    purpose: str = """
        Collect information about a Commercial Lease.
        """
    terms: Lease
    result: str = ""

    def handle(self) -> str:
        print(
            f"""
        DONE! Successfully extracted Lease Info:
        {self.terms}
        """
        )
        return "DONE " + json.dumps(self.terms.model_dump())

    @classmethod
    def format_instructions(cls, tool: bool = True) -> str:
        instr = super().format_instructions(tool)
        instr += """
        ------------------------------
        ASK ME QUESTIONS ONE BY ONE, to FILL IN THE FIELDS 
        of the `lease_info` function/tool.
        First ask me for the start date of the lease.
        DO NOT ASK ANYTHING ELSE UNTIL YOU RECEIVE MY ANSWER.
        """
        return instr

    @classmethod
    def examples(cls) -> List["LeaseMessage"]:
        return [
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-01-01", end_date="2021-12-31"),
                    financials=LeaseFinancials(monthly_rent="$1000", deposit="$1000"),
                    address="123 Main St, San Francisco, CA 94105",
                ),
                result="",
            ),
        ]


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    llm_cfg = OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust based on model
        temperature=0,
        timeout=45,
    )
    doc_agent = DocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            n_neighbor_chunks=2,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=50,
                overlap=10,
                n_neighbor_ids=4,
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.vecdb.set_collection("docqa-chat-multi-extract", replace=True)
    print("[blue]Welcome to the real-estate info-extractor!")
    doc_agent.config.doc_paths = [
        "examples/docqa/lease.txt",
    ]
    doc_agent.ingest()
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        system_message="""You are an expert on Commercial Leases. 
        You will receive various questions about a Commercial 
        Lease contract, along with some excerpts from the Lease.
        Your job is to answer them concisely in at most 2 sentences.
        """,
    )

    lease_extractor_agent = ChatAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    lease_extractor_agent.enable_message(LeaseMessage)

    lease_task = Task(
        lease_extractor_agent,
        name="LeaseExtractorAgent",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message=f"""
        You are an expert at understanding JSON function/tool specifications, and
        you are also very familiar with commercial lease terminology and concepts.
         
        See the `lease_info` function/tool below,  Your FINAL GOAL is to fill
        in the required fields in this `lease_info` function/tool,
        as shown in the example. This is ONLY an EXAMPLE,
        and YOU CANNOT MAKE UP VALUES FOR THESE FIELDS.
        
        To fill in these fields, you must ASK ME QUESTIONS about the lease,
        ONE BY ONE, and I will answer each question. 
        If I am unable to answer your question initially, try asking me 
        differently. If I am still unable to answer after 3 tries, fill in 
        {NO_ANSWER} for that field.
        When you have collected this info, present it to me using the 
        'lease_info' function/tool.
        DO NOT USE THIS Function/tool UNTIL YOU HAVE ASKED QUESTIONS 
        TO FILL IN ALL THE FIELDS.
        
        Think step by step. 
        Phrase each question simply as "What is ... ?",
        and do not explain yourself, or say any extraneous things. 
        Start by asking me for the start date of the lease.
        When you receive the answer, then ask for the next field, and so on.
        """,
    )

    validator_agent = ChatAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
            system_message=f"""
            You are obedient, understand instructions, and follow them faithfully,
            paying attention to the FORMAT specified,
            and you are also extremely CONCISE and SUCCINCT in your responses.
            
            Your task is to detect if the user's message is asking for ONE
            piece of information, or MULTIPLE pieces. Here is how you respond:
            
            IF the msg is only asking about ONE thing, OR if it is NOT EVEN a question:
                respond '{DONE}' and say nothing else.

            IF the msg is asking MORE THAN ONE thing,  respond like this:
            "Please only ask ONE question at a time. Try your question again.
            ONLY when you have ALL the answers, then present the info
            using the `lease_info` function/tool."
            """,
        )
    )
    validator_task = Task(
        validator_agent,
        name="Validator",
        single_round=True,
        interactive=False,
    )

    lease_task.add_sub_task([validator_task, doc_task])
    lease_task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/chat-multi-extract-local.py">
"""
Extract structured info from a commercial lease document,
using multiple agents, powered by a weaker/local LLM, combining tools/functions and RAG.

TASK:
Given a lease document, generate the lease terms, organized into
 a nested JSON structure defined by the Pydantic class `Lease`

Solution with Langroid Agents and tools:
1. QuestionGeneratorAgent: Lease JSON Spec -> list of questions to ask
2. InterrogatorAgent: For each question, generate 2 variants of the question,
   so we use total 3 variants per question, joined together, to increase
   the likelihood of getting an answer from the DocAgent (RAG).
3. DocAgent (has access to the lease) -> answer one question using RAG
3. LeasePresenterAgent: List of (question, answer) pairs ->
        organized into specified Lease JSON structure

Run like this:
```
python3 examples/docqa/chat-multi-extract-local.py -m ollama/mistral:7b-instruct-v0.2-q8_0
```
This works with a local mistral-instruct-v0.2 model.
(To use with ollama, first do `ollama run <model>` then
specify the model name as -m ollama/<model>)

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

Optional script args:
-m <local-model-name>, e.g. -m ollama/mistral:7b-instruct-v0.2-q8_0
(if omitted, defaults to GPT4o)
-nc to disable cache retrieval
-d to enable debug mode: see prompts, agent msgs etc.
"""

import json
import os
from typing import List, Optional

import typer
from rich import print

import langroid.language_models as lm
from langroid.agent import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE, NO_ANSWER
from langroid.utils.pydantic_utils import get_field_names

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class LeasePeriod(BaseModel):
    start_date: str
    end_date: str


class LeaseFinancials(BaseModel):
    monthly_rent: str
    deposit: str


class Lease(BaseModel):
    """
    Various lease terms.
    Nested fields to make this more interesting/realistic
    """

    period: LeasePeriod
    financials: LeaseFinancials
    address: str


class QuestionsTool(ToolMessage):
    request: str = "questions_tool"
    purpose: str = """
    To present a list of <questions> to ask, to fill a desired JSON structure.
    """
    questions: List[str]


class QuestionGeneratorAgent(ChatAgent):
    questions_list: List[str] = []

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
            return """
            You forgot to present the information in JSON format 
            according to the `questions_tool` specification,
            or you may have used a wrong tool name or field name.
            Remember that you must include `request` and `questions` fields,
            where `request` is "questions_tool" and `questions` is a list of questions.
            Try again.
            """
        return None

    def questions_tool(self, msg: QuestionsTool) -> str:
        # get all the field names, including nested ones
        fields = get_field_names(Lease)
        if len(msg.questions) < len(fields):
            return f"""
            ERROR: Expected {len(fields)} questions, but only got {len(msg.questions)}.
            See what you may have missed and try again.
            Hint: the required fields are {fields}
            """
        elif len(msg.questions) > len(fields):
            return f"""
            ERROR: Expected {len(fields)} questions, but got {len(msg.questions)}.
            You generated an extra question. Try again.
            Hint: the required fields are {fields}
            """
        else:
            self.questions_list = msg.questions
            return DONE + json.dumps(msg.questions)


class MyDocChatAgent(DocChatAgent):
    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        """
        Override the default LLM response to return the full document,
        to forget the last round in conversation, so we don't clutter
        the chat history with all previous questions
        (Assume questions don't depend on past ones, as is the case here,
        since we are extracting separate pieces of info from docs)
        """
        n_msgs = len(self.message_history)
        response = super().llm_response(message)
        # If there is a response, then we will have two additional
        # messages in the message history, i.e. the user message and the
        # assistant response. We want to (carefully) remove these two messages.
        self.message_history.pop() if len(self.message_history) > n_msgs else None
        self.message_history.pop() if len(self.message_history) > n_msgs else None
        return response


class LeasePresenterAgent(ChatAgent):
    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        """Handle scenario where Agent failed to present the Lease JSON"""
        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
            return """
            You either forgot to present the information in the JSON format
            required in `lease_info` JSON specification,
            or you may have used the wrong name of the tool or fields.
            Try again.
            """
        return None


class LeaseMessage(ToolMessage):
    """Tool/function to use to present details about a commercial lease"""

    request: str = "lease_info"
    purpose: str = "To present the <terms> of a Commercial lease."
    terms: Lease

    def handle(self) -> str:
        print(
            f"""
        DONE! Successfully extracted Lease Info:
        {self.terms}
        """
        )
        return DONE + " " + json.dumps(self.terms.model_dump())

    @classmethod
    def examples(cls) -> List["LeaseMessage"]:
        return [
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-01-01", end_date="2021-12-31"),
                    financials=LeaseFinancials(monthly_rent="$1000", deposit="$1000"),
                    address="123 Main St, San Francisco, CA 94105",
                ),
                result="",
            ),
            cls(
                terms=Lease(
                    period=LeasePeriod(start_date="2021-04-01", end_date="2022-04-28"),
                    financials=LeaseFinancials(monthly_rent="$2000", deposit="$2000"),
                    address="456 Main St, San Francisco, CA 94111",
                ),
                result="",
            ),
        ]


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            cache_type="fakeredis",
        )
    )
    llm_cfg = OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=32_000,  # adjust based on model
        timeout=120,
        temperature=0.2,
    )

    # (1) QUESTION GENERATOR
    question_generator_agent = QuestionGeneratorAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
            system_message="""
            See the `lease_info` JSON structure below. 
            Your ONLY task is to generate 
            QUESTIONS corresponding to each field in the `lease_info` JSON,
            and present these to me using the `questions_tool` in JSON format.
            Pay attention to the format and fields in the `questions_tool` JSON.
            """,
        )
    )
    question_generator_agent.enable_message(LeaseMessage)
    question_generator_agent.enable_message(QuestionsTool)
    question_generator_task = Task(
        question_generator_agent,
        name="QuestionGeneratorAgent",
        interactive=False,
    )

    # (2) RAG AGENT: try to answer a given question based on documents
    doc_agent = MyDocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            assistant_mode=True,
            n_neighbor_chunks=2,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=150,
                overlap=30,
                n_neighbor_ids=4,
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.vecdb.set_collection("docqa-chat-multi-extract", replace=True)
    doc_agent.ingest_doc_paths(["examples/docqa/lease.txt"])
    print("[blue]Welcome to the real-estate info-extractor!")
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        interactive=False,
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        system_message="""You are an expert on Commercial Leases. 
        You will receive a question about a Commercial 
        Lease contract, and your job is to answer concisely in at most 2 sentences.
        """,
    )

    # (3) Interrogator: persists in getting an answer for a SINGLE question
    #       from the RAG agent
    interrogator = ChatAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
            system_message="""
            You are an expert on Commercial leases and their terms. 
            User will send you a QUESTION about such a lease.
            Your ONLY job is to reply with TWO VARIATIONS of the QUESTION,
            and say NOTHING ELSE.
            """,
        )
    )
    interrogator_task = Task(
        interrogator,
        name="Interrogator",
        restart=True,  # clear agent msg history
        interactive=False,
        single_round=True,
    )

    # (4) LEASE PRESENTER: Given full list of question-answer pairs,
    #       organize them into the Lease JSON structure
    lease_presenter = LeasePresenterAgent(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    lease_presenter.enable_message(LeaseMessage)

    lease_presenter_task = Task(
        lease_presenter,
        name="LeasePresenter",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message="""
        The user will give you a list of Questions and Answers 
        about a commercial lease.
        
        Organize this information into the `lease_info` JSON structure specified below,
        and present it to me. 
        For fields where the answer is NOT KNOWN, fill in "UNKNOWN" as the value.
        """,
    )

    # (5) Use the agents/tasks

    # Lease info JSON -> Questions
    question_generator_task.run()
    questions = question_generator_agent.questions_list
    print(f"found {len(questions)} questions! Now generating answers...")

    # Questions -> Answers using RAG
    answers = []
    for q in questions:
        # use 3 variants of the question at the same time,
        # to increase likelihood of getting an answer
        q_variants = interrogator_task.run(q).content
        result = doc_task.run(q + "\n" + q_variants)
        answer = result.content or NO_ANSWER
        answers.append(answer)
    print(f"got {len(answers)} answers!")

    q2a = dict(zip(questions, answers))
    print(f"q2a: {q2a}")
    questions_answers = "\n\n".join(
        f"Question: {q}:\nAnswer: {a}" for q, a in q2a.items()
    )
    # Questions + Answers -> organized into nested Lease Info JSON
    lease_presenter_task.run(questions_answers)


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/chat-qa-summarize.py">
"""
Two-agent system to do Question-Answer based summarization of documents.
E.g. one could use this to summarize a very large document, assuming there is a
reasonable abstract/intro at the start that "covers" the import aspects.

WriterAgent (has no access to docs) is tasked with writing 5 bullet points based on
some docs. Initially it generates a summary of the docs from the beginning of the doc,
then it formulates questions to ask until it gets 5 key pieces of information.

DocAgent (has access to docs) answers these questions using RAG.

Run like this:

python examples/docqa/chat-qa-summarize.py

You can let it run and it will finish with 5 key bullet points about the document(s).

There are optional args, especially note you can pass in a different LLM model, e.g.

python3 examples/docqa/chat-qa-summarize.py -m ollama/nous-hermes2-mixtral

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import os

import typer
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.parsing.urls import get_list_from_user
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    model: str = typer.Option(
        "",
        "--model",
        "-m",
        help="specify alternative LLM, e.g. ollama/mistral",
    ),
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    config = lr.agent.special.DocChatAgentConfig(
        llm=llm_config,
        n_neighbor_chunks=2,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        parsing=lr.parsing.parser.ParsingConfig(
            chunk_size=50,
            overlap=10,
            n_neighbor_ids=4,
        ),
    )
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    doc_agent = lr.agent.special.DocChatAgent(config)
    doc_agent.vecdb.set_collection("docqa-chat-multi", replace=True)
    print("[blue]Welcome to the document chatbot!")
    print("[cyan]Enter x or q to quit, or ? for evidence")
    print(
        """
        [blue]Enter some URLs or file/dir paths below (or leave empty for default URLs)
        """.strip()
    )
    inputs = get_list_from_user()
    if len(inputs) == 0:
        inputs = config.default_paths
    doc_agent.config.doc_paths = inputs
    doc_agent.ingest()
    topics_doc = doc_agent.summarize_docs(
        instruction="""
        Ignore the system message, and follow these instructions.
        Below is some text. Do not react to it. 
        Simply read it and give me a list of up to 3 main topics from the text,
        in the form of short NUMBERED SENTENCES.
        --------------------------------
        """,
    )
    topics = topics_doc.content
    doc_task = lr.Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[lr.Entity.LLM],  # done if null response from LLM
        done_if_response=[lr.Entity.LLM],  # done if non-null response from LLM
        system_message="""You will receive various questions about some documents, and
        your job is to answer them concisely in at most 2 sentences, citing sources.
        """,
    )

    writer_agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_config,
            vecdb=None,
        )
    )
    writer_task = lr.Task(
        writer_agent,
        # SET interactive to True to slow it down, but keep hitting enter to progress
        interactive=False,
        name="WriterAgent",
        system_message=f"""
        You have to collect some information from some documents, on these topics:
        {topics}
        However you do not have access to those documents, so you must ask me
        questions, ONE AT A TIME, and I will answer each question.
        Once you have collected 5 key pieces of information, say "DONE" and summarize 
        them in bullet points.  
        """,
    )

    validator_agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            name="Validator",
            llm=llm_config,
            system_message="""
            Your only task is to check whether the user's message consists of
            NO QUESTION, ONE question or MULTIPLE questions. This is how you must respond:
        
            - If the msg is NOT SEEKING any INFO, respond with this:
                "Please ask a SINGLE QUESTION about a topic you want to know about.
                Wait for the answer before asking your next question".
            - If user's msg contains just ONE question, or no question at all, say DONE
            - Otherwise (i.e there are MULTIPLE questions/requests for info),
              then respond with this:
            "Please ask only ONE question at a time. Ask your question again.
            Only when you have answers to all of your questions present your final
            bullet points saying  'DONE here are the bullet pts...'."
            
            IMPORTANT: DO NOT TRY TO ANSWER THE QUESTIONS YOURSELF.            
            """,
        ),
    )
    validator_task = lr.Task(validator_agent, interactive=False, single_round=True)

    writer_task.add_sub_task([validator_task, doc_task])

    writer_task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/chat-search-filter.py">
"""
Variant of chat-search.py that uses a filter to identify different
set of ingested docs (obtained from web-search), so that cross-doc
questions can be answered.

This is a single-agent question-answering system that has access to a Web-Search
Tool when needed,
and in case a web search is used, ingests scraped link contents into a vector-db,
and uses Retrieval Augmentation to answer the question.

Run like this:

    python3 examples/docqa/chat-search-filter.py

Optional args:
    -nc : turn off caching (i.e. don't retrieve cached LLM responses)
    -d: debug mode, to show all intermediate results
    -f: use OpenAI functions api instead of tools
    -m <model_name>:  (e.g. -m ollama/mistral:7b-instruct-v0.2-q4_K_M)
    (defaults to GPT4-Turbo if blank)

(See here for guide to using local LLMs with Langroid:)
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import json
import re
from typing import Any, List

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatDocument
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import ForwardTool
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.parsing.web_search import metaphor_search
from pydantic import Field
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER


class RelevantExtractsTool(ToolMessage):
    request: str = Field(
        "relevant_extracts", description="MUST be included in EVERY use of this tool!"
    )
    purpose: str = "Get docs/extracts relevant to the <query> from prior searches"
    query: str = Field(..., description="The query to get relevant extracts for")
    filter_tag: str = Field(
        "",
        description="""
        Optional LOWER-CASE tag to filter to use for the search, 
        to restrict relevance extraction to a SPECIFIC PRIOR search result.
        IMPORTANT - DO NOT INTRODUCE A NEW TAG HERE!! You MUST use ONLY a
        tag you previously used in the `relevant_search_extracts` tool,
        to correctly identify a prior search result.
        """,
    )

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(
                query="when was the Mistral LLM released?",
                filter_tags=["mistral", "llm"],
            ),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """


class RelevantSearchExtractsTool(ToolMessage):
    request: str = Field(
        "relevant_search_extracts",
        description="MUST be included in EVERY use of this tool!",
    )
    purpose: str = "Get docs/extracts relevant to the <query> from a web search"
    query: str = Field(..., description="The search query to get relevant extracts for")
    num_results: int = Field(3, description="The number of search results to use")
    tag: str = Field(
        "",
        description="""
        Optional LOWER-CASE tag to attach to the documents ingested from the search, 
        to UNIQUELY IDENTIFY the docs ingested from this search, for future reference
        when using the `relevant_extracts` tool.
        """,
    )

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(
                query="when was the Mistral LLM released?",
                num_results=3,
                tag="mistral",
            ),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        IMPORTANT: You must include an ACTUAL query in the `query` field,
        """


def tags_to_filter(tags: List[str]) -> str | None:
    """
    Given a list of tags, create a qdrant-db filter condition expressing:
    EVERY tag MUST appear in the metadata.tags field of the document.
    Args:
        tags: List of tags to filter by
    Returns:
        json string of the qdrant filter condition, or None
    """
    if len(tags) == 0:
        return None
    match_conditions = [
        {"key": "metadata.tags", "match": {"any": [tag]}} for tag in tags
    ]

    filter = {"must": match_conditions}
    return json.dumps(filter)


class SearchDocChatAgent(DocChatAgent):

    def init_state(self) -> None:
        super().init_state()
        self.original_docs = []
        self.tried_vecdb: bool = False

    def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
        if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
            # no tool, so it must be meant for user
            return ForwardTool(agent="user")

    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> ChatDocument | None:
        return ChatAgent.llm_response(self, message)

    def relevant_extracts(self, msg: RelevantExtractsTool) -> str:
        """Get docs/extracts relevant to the query, from vecdb"""
        self.tried_vecdb = True
        query = msg.query
        if msg.filter_tag != "":
            self.set_filter(tags_to_filter([msg.filter_tag]))
        _, extracts = self.get_relevant_extracts(query)
        if len(extracts) == 0:
            return """
            No extracts found! You can try doing a web search with the
            `relevant_search_extracts` tool/function-call.
            """
        return "\n".join(str(e) for e in extracts)

    def relevant_search_extracts(self, msg: RelevantSearchExtractsTool) -> str:
        """Get docs/extracts relevant to the query, from a web search"""
        # if not self.tried_vecdb and len(self.original_docs) > 0:
        #     return "Please try the `relevant_extracts` tool, before using this tool"
        self.tried_vecdb = False
        query = msg.query
        # if query contains a url, then no need to do web search --
        # just ingest the specific link in the query
        if "http" in query:
            # extract the URL from the query
            url = re.search(r"(?P<url>https?://[^\s]+)", query).group("url")
            links = [url]
            # remove the url from the query
            query = re.sub(r"http\S+", "", query)
        else:
            results = metaphor_search(query, msg.num_results)
            links = [r.link for r in results]
        self.ingest_doc_paths(links, metadata={"tags": [msg.tag]})
        if msg.tag != "":
            self.set_filter(tags_to_filter([msg.tag]))
        _, extracts = self.get_relevant_extracts(query)
        return "\n".join(str(e) for e in extracts)


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = "",
    fn_api: bool = True,
) -> None:

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    print(
        """
        [blue]Welcome to the Internet Search chatbot!
        I will try to answer your questions, relying on (full content of links from) 
        Duckduckgo (DDG) Search when needed.
        
        Enter x or q to quit, or ? for evidence
        """
    )

    system_msg = Prompt.ask(
        """
    [blue] Tell me who I am (give me a role) by completing this sentence: 
    You are...
    [or hit enter for default]
    [blue] Human
    """,
        default="a helpful assistant.",
    )
    system_msg = re.sub("you are", "", system_msg, flags=re.IGNORECASE)

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=2048,  # adjust based on model
    )

    config = DocChatAgentConfig(
        use_functions_api=fn_api,
        use_tools=not fn_api,
        llm=llm_config,
        extraction_granularity=3,
        # for relevance extraction
        # relevance_extractor_config=None,  # set to None to disable relevance extraction
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=200,  # aim for this many tokens per chunk
            overlap=50,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="fitz",
            ),
        ),
        system_message=f"""
        {system_msg} You will try your best to answer my questions,
        in this order of preference:
        1. If you can answer from your own knowledge, simply return the answer
        2. Otherwise:
         2.1 If the question contains a URL, then use the `relevant_search_extracts`
             tool/function with the `query` field set to 
             this EXACT QUESTION INTACT! (DO NOT REPHRASE IT),
             and set the appropriate `tag` to UNIQUELY identify 
             docs from this search, to be able to refer to docs from 
             this search in FUTURE uses of the `relevant_extracts` tool.
         2.1 Otherwise, 
             if you have previously used the `relevant_search_extracts` 
                tool/fn-call
             to do a web search, you can ask for some relevant text from those search
             results, using the `relevant_extracts` tool/function-call, 
             and you MUST ONLY use a PREVIOUSLY used tag to correctly identify
             the prior search results to narrow down the search,
             and you will receive relevant extracts, if any.
             If you receive {NO_ANSWER}, it means no relevant extracts exist,
             and you can try the next step 2.2, using a web search.
             
         2.2 otherwise, i.e. you have NOT YET done a web search, you can use
             the `relevant_search_extracts` tool/function-call to search the web,
             MAKING SURE YOU SET a UNIQUE TAG (LOWER CASE, short word or 
             phrase) in the `tag` field, to UNIQUELY identify the docs from 
             this search, to be able to refer to them in a future use of 
             `relevant_extracts` tool.
             You will then receive relevant extracts from these search results, 
             if any. 
        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} 
        
        Remember these simple rules:
         (a) if a question contains a URL, simply use the `relevant_search_extracts`
                tool/function-call with the `query` field set to this EXACT QUESTION
         (b) else if you have ALREADY done a web-search 
         (using the `relevant_search_extracts` tool),
         you should FIRST try `relevant_extracts` to see if there are
         any relevant passages from PREVIOUS SEARCHES, before doing a new search.
         
         YOU CAN USE TOOLS MULTIPLE TIMES before composing your answer.
         For example, when asked to compare two things, you can use the
         `relevant_extracts` tool multiple times to get relevant extracts
         from different PRIOR search results, and THEN compose your answer!
        
        Be very concise in your responses, use no more than 1-2 sentences.
        When you answer based on provided documents, be sure to show me 
        the SOURCE(s) and EXTRACT(s), for example:
        
        SOURCE: https://www.wikihow.com/Be-a-Good-Assistant-Manager
        EXTRACT: Be a Good Assistant ... requires good leadership skills.
        
        For the EXTRACT, ONLY show up to first 3 words, and last 3 words.
        """,
    )

    agent = SearchDocChatAgent(config)
    agent.enable_message(RelevantExtractsTool)
    agent.enable_message(RelevantSearchExtractsTool)
    collection_name = Prompt.ask(
        "Name a collection to use",
        default="docqa-chat-search",
    )
    replace = (
        Prompt.ask(
            "Would you like to replace (i.e. erase) this collection?",
            choices=["y", "n"],
            default="n",
        )
        == "y"
    )

    print(f"[red]Using {collection_name}")

    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?")


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/chat.py">
"""
Single agent to use to chat with a Retrieval-augmented LLM.
Repeat: User asks question -> LLM answers.

Run like this, either with a document-path (can be URL, file-path, folder-path):

python3 examples/docqa/chat.py url-or-file-orfolder-path

(or run with no arguments to go through the dialog).

If a document-arg is provided, it will be ingested into the vector database.

To change the model, use the --model flag, e.g.:

python3 examples/docqa/chat.py --model ollama/mistral:7b-instruct-v0.2-q8_0

To change the embedding service provider, use the --embed and --embedconfig flags, e.g.:

For OpenAI
python3 examples/docqa/chat.py --embed openai

For Huggingface SentenceTransformers
python3 examples/docqa/chat.py --embed hf --embedconfig BAAI/bge-large-en-v1.5

For Llama.cpp Server
python3 examples/docqa/chat.py --embed llamacpp --embedconfig localhost:8000

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import os

import typer
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


@app.command()
def main(
    doc: str = typer.Argument("", help="url, file-path or folder to chat about"),
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    vecdb: str = typer.Option(
        "qdrant", "--vecdb", "-v", help="vector db name (default: qdrant)"
    ),
    nostream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    embed_provider: str = typer.Option(
        "openai",
        "--embed",
        "-e",
        help="Embedding service provider",
        # openai, hf, llamacpp
    ),
    embed_config: str = typer.Option(
        None,
        "--embedconfig",
        "-ec",
        help="Embedding service host/sentence transformer model",
    ),
    # e.g. NeuML/pubmedbert-base-embeddings
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust as needed
        temperature=0.2,
        max_output_tokens=300,
        timeout=60,
    )

    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
        extraction_granularity=3,
        # for relevance extraction
        # relevance_extractor_config=None,  # set to None to disable relevance extraction
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=200,  # aim for this many tokens per chunk
            overlap=50,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # see here for possible values:
                # https://github.com/langroid/langroid/blob/main/langroid/parsing/parser.py
                library="pymupdf4llm",
            ),
        ),
    )

    match embed_provider:
        case "hf":
            embed_cfg = lr.embedding_models.SentenceTransformerEmbeddingsConfig(
                model_type="sentence-transformer",
                model_name=embed_config,
            )
        case "llamacpp":
            embed_cfg = lr.embedding_models.LlamaCppServerEmbeddingsConfig(
                api_base=embed_config,
                dims=768,  # Change this to match the dimensions of your embedding model
            )
        case "gemini":
            embed_cfg = lr.embedding_models.GeminiEmbeddingsConfig(
                model_type="gemini", dims=768
            )
        case _:
            embed_cfg = lr.embedding_models.OpenAIEmbeddingsConfig()

    match vecdb:
        case "lance" | "lancedb":
            config.vecdb = lr.vector_store.LanceDBConfig(
                collection_name="doc-chat-lancedb",
                storage_path=".lancedb/data/",
                embedding=embed_cfg,
            )
        case "qdrant" | "qdrantdb":
            config.vecdb = lr.vector_store.QdrantDBConfig(
                cloud=False,
                storage_path=".qdrant/doc-chat",
                embedding=embed_cfg,
            )
        case "chroma" | "chromadb":
            config.vecdb = lr.vector_store.ChromaDBConfig(
                storage_path=".chroma/doc-chat",
                embedding=embed_cfg,
            )
        case "weaviate" | "weaviatedb":
            config.vecdb = lr.vector_store.WeaviateDBConfig(
                embedding=embed_cfg,
            )
        case "pinecone" | "pineconedb":
            config.vecdb = lr.vector_store.PineconeDBConfig(
                collection_name="doc-chat-pinecone-serverless",
                embedding=embed_cfg,
            )
        case "postgres" | "postgresdb":
            config.vecdb = lr.vector_store.PostgresDBConfig(
                embedding=embed_cfg, cloud=True
            )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not nostream,
        )
    )

    agent = DocChatAgent(config)
    print("[blue]Welcome to the document chatbot!")

    if doc:
        # TODO - could save time by checking whether we've already ingested this doc(s)
        agent.ingest_doc_paths([doc])
    else:
        agent.user_docs_ingest_dialog()

    print("[cyan]Enter x or q to quit")

    task = lr.Task(
        agent,
        system_message="You are a helpful assistant, "
        "answering questions about some docs",
    )
    task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/crawl4ai_examples.py">
import json
import os
from typing import Optional

from langroid.parsing.url_loader import Crawl4aiConfig, URLLoader
from crawl4ai.async_configs import LLMConfig
from crawl4ai.extraction_strategy import (
    JsonCssExtractionStrategy,
    LLMExtractionStrategy,
)
from crawl4ai.markdown_generation_strategy import DefaultMarkdownGenerator
from crawl4ai.content_filter_strategy import (
    PruningContentFilter,
    LLMContentFilter,
)
from crawl4ai.content_scraping_strategy import LXMLWebScrapingStrategy

# Helper for pydantic models if LLMExtractionStrategy is used with schema
from pydantic import BaseModel, Field
from typing import List
from langroid.mytypes import Document


from rich.console import Console
from rich.prompt import IntPrompt

console = Console()

import sys

sys.path.append(os.path.dirname(__file__))


def simple_crawler_example():
    """
    Demonstrates a basic crawl using Crawl4aiConfig with default settings.
    It will fetch the markdown content of the given URLs.
    """
    print("\n--- Running simple_crawler_example ---")
    urls = [
        "https://pytorch.org",
        "https://arxiv.org/pdf/1706.03762",  # This will be handled by DocumentParser
    ]
    crawler_config = Crawl4aiConfig()  # Uses default BrowserConfig and CrawlerRunConfig
    loader = URLLoader(urls=urls, crawler_config=crawler_config)

    docs = loader.load()
    for doc in docs:
        print(
            f"URL: {doc.metadata.source}, Content Length: {len(doc.content)} (first 200 chars: {doc.content[:200]})"
        )
    print("--- simple_crawler_example finished ---")


def extract_to_json_example():
    """
    Demonstrates how to use `JsonCssExtractionStrategy` to extract structured JSON
    from a webpage, configured via `Crawl4aiConfig`.
    """
    print("\n--- Running extract_to_json_example ---")
    HACKER_NEWS_URL = "https://news.ycombinator.com"
    HACKER_NEWS_SCHEMA = {
        "name": "HackerNewsArticles",
        "baseSelector": "tr.athing",  # Each article is in a <tr> with class 'athing'
        "fields": [
            {"name": "title", "selector": "span.titleline > a", "type": "text"},
            {
                "name": "link",
                "selector": "span.titleline > a",
                "type": "attribute",
                "attribute": "href",
            },
        ],
    }

    css_strategy = JsonCssExtractionStrategy(schema=HACKER_NEWS_SCHEMA)

    hn_crawler_config = Crawl4aiConfig(extraction_strategy=css_strategy)

    print(f"Starting scrape of {HACKER_NEWS_URL}...")
    loader = URLLoader(urls=[HACKER_NEWS_URL], crawler_config=hn_crawler_config)
    documents = loader.load()

    if documents:
        print("\nScrape successful! Processing extracted data...")
        extracted_json_string = documents[0].content
        try:
            extracted_data = json.loads(extracted_json_string)
            print("\n--- Top 3 Articles from Hacker News ---")
            for i, item in enumerate(extracted_data[:3], 1):
                print(f"{i}. Title: {item.get('title')}")
                print(f"   Link: {item.get('link')}")
            print(f"\nTotal items extracted: {len(extracted_data)}")
        except json.JSONDecodeError:
            print("Error: Failed to parse the extracted content as JSON.")
            print("Received content:", extracted_json_string)
    else:
        print("\nScrape failed. No documents were returned.")
    print("--- extract_to_json_example finished ---")


def markdown_generation_example():
    """
    Demonstrates customizing markdown generation using `markdown_strategy` in Crawl4aiConfig.
    Uses PruningContentFilter for focused content.
    """
    print("\n--- Running markdown_generation_example ---")
    url = "https://news.ycombinator.com"

    # Define a content filter to prune irrelevant sections
    prune_filter = PruningContentFilter(
        threshold=0.6,  # More aggressive pruning
        threshold_type="dynamic",
        min_word_threshold=10,
    )

    # Configure the markdown generator to use the filter and ignore links
    md_generator = DefaultMarkdownGenerator(
        content_filter=prune_filter,
        options={
            "ignore_links": True,
            "body_width": 100,  # Wrap text at 100 characters
            "citations": False,  # Disable citations
        },
    )

    crawler_config = Crawl4aiConfig(markdown_strategy=md_generator)

    loader = URLLoader(urls=[url], crawler_config=crawler_config)
    docs = loader.load()

    if docs:
        print(f"Markdown Content (first 500 chars) for {url}:")
        # In this setup, the 'content' of the Document will be the fit_markdown
        print(docs[0].content[:500])
        print(f"Original URL: {docs[0].metadata.source}")
    else:
        print(f"Failed to crawl {url}.")
    print("--- markdown_generation_example finished ---")


def deep_crawl_example():
    """Crawl multiple pages from a domain using BFS strategy."""
    from crawl4ai.deep_crawling import BFSDeepCrawlStrategy
    from crawl4ai.deep_crawling.filters import (
        FilterChain,
        URLPatternFilter,
        DomainFilter,
        ContentTypeFilter,
    )
    from crawl4ai.async_configs import BrowserConfig, CrawlerRunConfig

    # Create browser config
    browser_config = BrowserConfig(
        # Example browser config settings
        headless=True,
        viewport={"width": 1920, "height": 1080},
    )

    # Create filter chain
    filter_chain = FilterChain(
        [
            URLPatternFilter(patterns=["*core*"]),
            DomainFilter(
                allowed_domains=["docs.crawl4ai.com"],
            ),
            ContentTypeFilter(allowed_types=["text/html"]),
        ]
    )

    # Create deep crawl strategy
    deep_crawl_strategy = BFSDeepCrawlStrategy(
        max_depth=2, include_external=False, max_pages=5, filter_chain=filter_chain
    )

    # Create run config
    run_config = CrawlerRunConfig(
        # Example run config settings
        deep_crawl_strategy=deep_crawl_strategy,
    )

    # Create the Crawl4ai configuration with all components
    crawler_config = Crawl4aiConfig(
        crawl_mode="deep", browser_config=browser_config, run_config=run_config
    )

    url = "https://docs.crawl4ai.com/"

    loader = URLLoader(urls=[url], crawler_config=crawler_config)

    docs = loader.load()

    if docs:
        print(f"Total Documents: {len(docs)}")
        for i, doc in enumerate(docs[:5], 1):
            print(f"{i}. {doc.metadata.source} ({len(doc.content)} chars)")
    else:
        print("No documents crawled.")


def scraping_strategy_example():
    """
    Demonstrates using a custom `scraping_strategy` (e.g., LXMLWebScrapingStrategy)
    in Crawl4aiConfig for potentially faster HTML parsing.
    """
    print("\n--- Running scraping_strategy_example ---")
    url = "https://www.nbcnews.com/business"

    # Use LXMLWebScrapingStrategy for potentially faster scraping
    scraping_strategy = LXMLWebScrapingStrategy()

    crawler_config = Crawl4aiConfig(scraping_strategy=scraping_strategy)

    print(f"Starting crawl of {url} with LXML scraping strategy...")
    loader = URLLoader(urls=[url], crawler_config=crawler_config)
    docs = loader.load()

    if docs:
        print(f"Crawl successful! Content Length for {url}: {len(docs[0].content)}")
        print(f"First 200 chars of content:\n{docs[0].content[:200]}")
    else:
        print(f"Failed to crawl {url}.")
    print("--- scraping_strategy_example finished ---")


def llm_extraction_example():
    """
    Demonstrates using LLMExtractionStrategy to extract structured data
    using an LLM, configured via Crawl4aiConfig.
    Requires GEMINI_API_KEY environment variable to be set.
    """
    print("\n--- Running llm_extraction_example ---")

    if not os.getenv("GEMINI_API_KEY"):
        print("GEMINI_API_KEY not found. Skipping llm_extraction_example.")
        print("Please set the GEMINI_API_KEY environment variable to run this example.")
        return

    class ArticleData(BaseModel):
        headline: str
        summary: str = Field(description="A short summary of the article")
        author: Optional[str] = None

    url = "https://news.ycombinator.com"

    llm_strategy = LLMExtractionStrategy(
        llm_config=LLMConfig(
            # Corrected Gemini model name based on your provided documentation
            provider="gemini/gemini-2.0-flash",
            api_token=os.getenv("GEMINI_API_KEY"),
        ),
        schema=ArticleData.schema_json(),
        extraction_type="schema",
        instruction="Extract the headline and a short summary for the main article on the page. If author is available, extract it too.",
        # Small chunk_token_threshold for demo purposes, adjust as needed for full pages
        chunk_token_threshold=1000,
        apply_chunking=True,
        input_format="markdown",  # Can be "html", "fit_markdown"
    )

    crawler_config = Crawl4aiConfig(extraction_strategy=llm_strategy)

    print(f"Starting LLM-based extraction from {url}...")
    loader = URLLoader(urls=[url], crawler_config=crawler_config)
    docs: List[Document] = loader.load()  # Explicitly type hint for clarity

    # The output structure is `[Document(...)]` because URLLoader wraps the result.
    # The actual extracted JSON is in `docs[0].content`.
    print(
        f"Raw documents loaded: {docs}"
    )  # This will show the `Document` object structure

    if docs:
        print("\nLLM Extraction successful!")
        extracted_content = docs[0].content
        try:
            # LLM extraction returns JSON string in `content`
            extracted_data = json.loads(extracted_content)
            print("Extracted Data:", json.dumps(extracted_data, indent=2))
        except json.JSONDecodeError as e:
            print(f"Error parsing LLM output JSON: {e}")
            print("Raw LLM output:", extracted_content)
    else:
        print(f"LLM extraction from {url} failed or returned no data.")
    print("--- llm_extraction_example finished ---")


def regex_extraction_example():
    """
    Demonstrates using RegexExtractionStrategy to extract URLs, emails, and dates
    from a webpage, configured via Crawl4aiConfig.
    """
    from langroid.parsing.url_loader import Crawl4aiConfig, URLLoader
    from crawl4ai.extraction_strategy import RegexExtractionStrategy
    from langroid.mytypes import Document
    import json

    print("\n--- Running regex_extraction_example ---")

    # Pick a real-world page that likely has email, URL, or date patterns
    url = "https://www.scrapethissite.com/pages/forms/"

    # Combine multiple regex types
    regex_strategy = RegexExtractionStrategy(
        pattern=(
            RegexExtractionStrategy.Email
            | RegexExtractionStrategy.Url
            | RegexExtractionStrategy.DateUS
        ),
    )

    crawler_config = Crawl4aiConfig(extraction_strategy=regex_strategy)

    print(f"Crawling and extracting from: {url}")
    loader = URLLoader(urls=[url], crawler_config=crawler_config)
    docs = loader.load()

    if not docs:
        print("No documents returned.")
        return

    try:
        extracted_json = json.loads(docs[0].content)
        if not isinstance(extracted_json, list) or not extracted_json:
            print("No structured matches found.")
            return

        print(f"Found {len(extracted_json)} matches:")
        for i, item in enumerate(extracted_json[:10], start=1):  # Show top 10
            label = item.get("label", "unknown")
            value = item.get("value", "")
            print(f"  {i}. [{label}] {value}")
    except json.JSONDecodeError:
        print("Failed to parse content as JSON.")
        print("Raw content:")
        print(docs[0].content)

    print("--- regex_extraction_example finished ---")


def llm_content_filter_example():
    """
    Demonstrates using LLMContentFilter within DefaultMarkdownGenerator
    to intelligently filter and format content.
    Requires GEMINI_API_KEY environment variable.
    """
    print("\n--- Running llm_content_filter_example ---")

    if not os.getenv("GEMINI_API_KEY"):
        print("GEMINI_API_KEY not found. Skipping llm_content_filter_example.")
        print("Please set the GEMINI_API_KEY environment variable to run this example.")
        return

    url = "https://news.ycombinator.com"  # A page with varied content

    llm_filter = LLMContentFilter(
        llm_config=LLMConfig(
            provider="gemini/gemini-2.0-flash",
            api_token=os.getenv("GEMINI_API_KEY"),
        ),
        instruction="""
        Focus on extracting the core news headlines and summaries.
        Include:
        - Main headlines
        - Brief summaries of the linked articles (if visible on the page)
        Exclude:
        - Navigation elements, sidebars, footer content
        - Comments sections
        Format the output as clean markdown with proper code blocks and headers if applicable.
        """,
        chunk_token_threshold=2048,  # Adjust for performance/cost
        verbose=False,  # Set to True for detailed LLM logs
    )

    md_generator = DefaultMarkdownGenerator(content_filter=llm_filter)

    crawler_config = Crawl4aiConfig(markdown_strategy=md_generator)

    print(f"Starting crawl of {url} with LLM content filter...")
    loader = URLLoader(urls=[url], crawler_config=crawler_config)
    docs = loader.load()

    if docs:
        print("\nLLM Content Filter successful!")
        # The content of the Document will be the `fit_markdown` from the LLM filter
        print("Filtered Markdown (first 1000 chars):")
        print(docs[0].content[:1000])
    else:
        print(f"LLM content filter crawl from {url} failed or returned no data.")
    print("--- llm_content_filter_example finished ---")


example_functions = {
    1: ("Simple Crawl Example", simple_crawler_example),
    2: ("JSON Extraction via CSS Selectors", extract_to_json_example),
    3: ("Custom Markdown Generation", markdown_generation_example),
    4: ("Deep Crawl with BFS", deep_crawl_example),
    5: ("LXML Scraping Strategy", scraping_strategy_example),
    6: ("LLM-based Extraction", llm_extraction_example),
    7: ("Regex Extraction", regex_extraction_example),
    8: ("LLM Content Filter in Markdown", llm_content_filter_example),
}


def main_menu():
    console.rule("[bold green]Crawl4ai Example Menu")
    for i, (name, _) in example_functions.items():
        console.print(f"[cyan]{i}.[/cyan] {name}")
    console.print("[magenta]0.[/magenta] Exit")

    while True:
        try:
            choice = IntPrompt.ask("\nChoose an example to run", default=0)

            if choice == 0:
                console.print("[bold red]Exiting. Goodbye![/bold red]")
                break
            elif choice in example_functions:
                console.rule(f"[bold yellow]Running: {example_functions[choice][0]}")
                example_functions[choice][1]()  # Run selected function
                console.print("\n[green] Finished.[/green]\n")
            else:
                console.print("[red]Invalid choice. Try again.[/red]")
        except KeyboardInterrupt:
            console.print("\n[bold red]Interrupted by user. Exiting.[/bold red]")
            break


if __name__ == "__main__":
    main_menu()
</file>

<file path="examples/docqa/doc-aware-chat.py">
"""
Single Agent for Doc-aware chat with user.

- user asks question
- LLM decides whether to:
    - ask user for follow-up/clarifying information, or
    - retrieve relevant passages from documents, or
    - provide a final answer, if it has enough information from user and documents.

To reduce response latency, in the DocChatAgentConfig,
you can set the `relevance_extractor_config=None`,
to turn off the relevance_extraction step, which uses the LLM
to extract verbatim relevant portions of retrieved chunks.

Run like this:

python3 examples/docqa/doc-aware-chat.py
"""

import os
from typing import Any, Optional

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.chat_agent import ChatAgent
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.agent.tools.orchestration import ForwardTool
from langroid.agent.tools.retrieval_tool import RetrievalTool
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class DocAwareChatAgent(DocChatAgent):
    def __init__(self, config: DocChatAgentConfig):
        super().__init__(config)
        self.enable_message(RetrievalTool)

    def retrieval_tool(self, msg: RetrievalTool) -> str:
        results = super().retrieval_tool(msg)
        return f"""
        
        RELEVANT PASSAGES:
        =====        
        {results}        
        ====
        
        
        BASED on these RELEVANT PASSAGES, DECIDE:
        - If this is sufficient to provide the user a final answer specific to 
            their situation, do so.
        - Otherwise, 
            - ASK the user for more information to get a better understanding
              of their situation or context, OR
            - use this tool again to get more relevant passages.
        """

    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        # override DocChatAgent's default llm_response
        return ChatAgent.llm_response(self, message)

    def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
        # we are here if there is no tool in the msg
        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
            # Any non-tool message must be meant for user, so forward it to user
            return ForwardTool(agent="User")


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = lm.OpenAIChatModel.GPT4o,
) -> None:
    llm_config = lm.OpenAIGPTConfig(chat_model=model)
    config = DocChatAgentConfig(
        llm=llm_config,
        n_query_rephrases=0,
        hypothetical_answer=False,
        relevance_extractor_config=None,
        # this turns off standalone-query reformulation; set to False to enable it.
        assistant_mode=True,
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=100,  # aim for this many tokens per chunk
            n_neighbor_ids=5,
            overlap=20,  # overlap between chunks
            max_chunks=10_000,
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="fitz",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    doc_agent = DocAwareChatAgent(config)
    print("[blue]Welcome to the document chatbot!")
    url = Prompt.ask("[blue]Enter the URL of a document")
    doc_agent.ingest_doc_paths([url])

    # For a more flexible/elaborate user doc-ingest dialog, use this:
    # doc_agent.user_docs_ingest_dialog()

    doc_task = Task(
        doc_agent,
        interactive=False,
        name="DocAgent",
        system_message=f"""
        You are a DOCUMENT-AWARE-GUIDE, but you do NOT have direct access to documents.
        Instead you can use the `retrieval_tool` to get passages from the documents
        that are relevant to a certain query or search phrase or topic.
        DO NOT ATTEMPT TO ANSWER THE USER'S QUESTION WITHOUT RETRIEVING RELEVANT
        PASSAGES FROM THE DOCUMENTS. DO NOT use your own existing knowledge!!
        Everything you tell the user MUST be based on the documents.
        
        The user will ask you a question that you will NOT be able to answer
        immediately, because you are MISSING some information about:
            - the user or their context or situation, etc
            - the documents relevant to the question
        
        At each turn you must decide among these possible ACTIONS:
        - use the `{RetrievalTool.name()}` to get more relevant passages from the 
            documents, OR
        - ANSWER the user if you think you have enough information 
            from the user AND the documents, to answer the question.
            
        You can use the `{RetrievalTool.name()}` multiple times to get more 
        relevant passages, if you think the previous ones were not sufficient.
        
        REMEMBER - your goal is to be VERY HELPFUL to the user; this means
        you should NOT OVERWHELM them by throwing them a lot of information and
        ask them to figure things out. Instead, you must GUIDE them 
        by asking SIMPLE QUESTIONS, ONE at at time, and finally provide them
        a clear, DIRECTLY RELEVANT answer that is specific to their situation. 
        """,
    )

    print("[cyan]Enter x or q to quit, or ? for evidence")

    doc_task.run("Can you help me with some questions?")


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/doc-aware-compose-2.py">
"""
2-agent doc-aware conversation,
different from standard question -> answer RAG flow.

Similar to doc-aware-guide-2.py, but in this case, the goal is
not to answer a user question, but to generate/compose a specific type of document,
adhering to some requirements, which are specified in external docs.

The DocAgent has access to the "requirement" docs.

To make this meaningful, ensure that the document(s) you upload
pertain to some type of "requirements" that a final generated document must adhere to.
E.g., try this document that specifies what a residential lease agreement should contain:
https://www.apartments.com/rental-manager/resources/leases/how-write-lease-agreement
And in the ensuing dialog, when prompted, say:
"I want to write an informal residential lease agreement."

ComposerAgent composes a document, via a multi-step
conversation, where it could either address:
- DocAgent (who has access to requirement docs) for info on requirements, or
- User, to ask follow-up questions about their situation/context.

python3 examples/docqa/doc-aware-guide-2.py

"""

import os
from typing import Optional

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import AT, DONE, NO_ANSWER

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class DocAgent(DocChatAgent):
    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        # Augment the response
        results = super().llm_response(message).content
        return self.create_llm_response(
            f"""
            Summary answer FROM DocAgent:
            
            ===
            {results}
            ===
            
            Look at the results above. These might be too much for the user to read.
            DECIDE whether you want to:
            - Ask the User a SINGLE follow-up question to get more info about their 
                situation or context, OR
            - Ask the DocAgent for more information, if you think you need more info.
            - Provide the User a FINAL answer, if you think you have enough information 
               from the User AND the Documents
               
            IMPORTANT: Do NOT simply give the User a list of options -- 
                you must HELP the user by asking them FOLLOWUP questions
                about their situation and GUIDE them to a SPECIFIC, 
                DIRECTLY RELEVANT answer. 
            REMEMBER - NEVER ask the DocAgent or User MULTIPLE questions at a time,
                always ask ONE question at a time.
            """
        )


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = lm.OpenAIChatModel.GPT4o,
) -> None:
    vecdb_config = lr.vector_store.QdrantDBConfig(
        storage_path=".qdrant/doc-aware/",
        replace_collection=False,
        cloud=False,
    )

    llm_config = lm.OpenAIGPTConfig(chat_model=model)
    config = DocChatAgentConfig(
        vecdb=vecdb_config,
        llm=llm_config,
        n_query_rephrases=0,
        hypothetical_answer=False,
        assistant_mode=True,
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=100,  # aim for this many tokens per chunk
            n_neighbor_ids=5,
            overlap=20,  # overlap between chunks
            max_chunks=10_000,
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    doc_agent = DocAgent(config)
    print("[blue]Welcome to the document chatbot!")
    doc_agent.user_docs_ingest_dialog()
    print("[cyan]Enter x or q to quit, or ? for evidence")
    doc_task = Task(
        doc_agent,
        interactive=False,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
    )

    composer = ChatAgent(
        ChatAgentConfig(
            name="Composer",
            llm=llm_config,
            vecdb=None,
        )
    )
    # MyRecipientTool = RecipientTool.create(
    #     recipients=["DocAgent", "User"], default="User"
    # )
    # guide_agent.enable_message(MyRecipientTool)
    task_config = lr.TaskConfig(addressing_prefix=AT)
    composer_task = Task(
        composer,
        interactive=False,
        config=task_config,
        system_message=f"""
        You are a SKILLFUL WRITER, who can adhere to specified REQUIREMENTS
        and GUIDELINES to generate specific types of documents. Your TASK
        is to create a certain type of Document requested by the User,
        IN MARKDOWN FORMAT,
        based on both the USER's info and the REQUIREMENTS specified in the
        specific DOCUMENTS.  
        
        However you do NOT have direct access to the specification docs, but you have an 
        assistant named DocAgent, who DOES have access to the documents.
          
        Since you could be talking to TWO people, in order to CLARIFY who you are
        addressing, you MUST ALWAYS EXPLICITLY ADDRESS either the 
        "User" or the "DocAgent" using {AT}User or {AT}DocAgent, respectively.
        
        You must THINK like this at each step after receiving a 
        DOCUMENT REQUEST from the User:
        
        (I NEVER WANT TO Overwhelm DocAgent or User with TOO MANY QUESTIONS,
        so I will ALWAYS ask ONE question at a time)
        
        - I must first find out more about the type of document the user
         wants, from DocAgent, let me address DocAgent to get the requirements info.
        - I got some info from DocAgent, let me now ask the User a follow-up question
            to get ONE SPECIFIC piece of information about their situation.
        - I need to get MORE info from DocAgent, let me ask DocAgent for more info.
        - DocAgent said {NO_ANSWER}!!, Let me try asking a different way.
        - I have a bit more info, now let me ask the User a further follow-up question,
            to get ONE SPECIFIC piece of information about their situation.
        - I need more info from user, let me ask the User a follow-up question,
            to get ANOTHER SPECIFIC piece of information about their situation.
        ...[and so on]...
        - Now I have ALL the info I need from BOTH the User and DocAgent,
            so I can provide the User the FINAL DOCUMENT, formatted 
             nicely in MARKDOWN, as per the requirements.
            so I will say {DONE}, followed by my composed document.   
            
        IMPORTANT: When giving the User a list of choices, always show them
            a NUMBERED list of choices.       
            
        ASK AT MOST 5 QUESTIONS TO THE USER, then generate the requested
        document to the best of your ability.   
        """,
    )
    composer_task.add_sub_task(doc_task)

    while True:
        query = Prompt.ask("[blue]How can I help?")
        if query in ["x", "q"]:
            break
        composer_task.run(query)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/doc-aware-guide-2.py">
"""
2-agent doc-aware conversation,
different from standard question -> answer RAG flow.

GuideAgent answers the user's question, via a multi-step
conversation, where it could either address:
- DocAgent (who has access to docs) for info, or
- User, to ask follow-up questions about their situation/context.

python3 examples/docqa/doc-aware-guide-2.py

"""

import os
from typing import Optional

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import AT, DONE, NO_ANSWER

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class DocAgent(DocChatAgent):
    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        # Augment the response
        results = super().llm_response(message).content
        return self.create_llm_response(
            f"""
            Summary answer FROM DocAgent:
            
            ===
            {results}
            ===
            
            Look at the results above. These might be too much for the user to read.
            DECIDE whether you want to:
            - Ask the User a SINGLE follow-up question to get more info about their 
                situation or context, OR
            - Ask the DocAgent for more information, if you think you need more info.
            - Provide the User a FINAL answer, if you think you have enough information 
               from the User AND the Documents
               
            IMPORTANT: Do NOT simply give the User a list of options -- 
                you must HELP the user by asking them FOLLOWUP questions
                about their situation and GUIDE them to a SPECIFIC, 
                DIRECTLY RELEVANT answer. 
            REMEMBER - NEVER ask the DocAgent or User MULTIPLE questions at a time,
                always ask ONE question at a time.
            """
        )


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = lm.OpenAIChatModel.GPT4o,
) -> None:
    vecdb_config = lr.vector_store.QdrantDBConfig(
        storage_path=".qdrant/doc-aware/",
        replace_collection=False,
        cloud=False,
    )

    llm_config = lm.OpenAIGPTConfig(chat_model=model)
    config = DocChatAgentConfig(
        vecdb=vecdb_config,
        llm=llm_config,
        n_query_rephrases=0,
        hypothetical_answer=False,
        assistant_mode=True,
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=100,  # aim for this many tokens per chunk
            n_neighbor_ids=5,
            overlap=20,  # overlap between chunks
            max_chunks=10_000,
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    doc_agent = DocAgent(config)
    print("[blue]Welcome to the document chatbot!")
    doc_agent.user_docs_ingest_dialog()
    print("[cyan]Enter x or q to quit, or ? for evidence")
    doc_task = Task(
        doc_agent,
        interactive=False,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
    )

    guide_agent = ChatAgent(
        ChatAgentConfig(
            name="GuideAgent",
            llm=llm_config,
            vecdb=None,
        )
    )
    # MyRecipientTool = RecipientTool.create(
    #     recipients=["DocAgent", "User"], default="User"
    # )
    # guide_agent.enable_message(MyRecipientTool)
    task_config = lr.TaskConfig(addressing_prefix=AT)
    guide_task = Task(
        guide_agent,
        interactive=False,
        config=task_config,
        system_message=f"""
        You are VERY HELPFUL GUIDE, who wants to help a User with their inquiry.
        
        Your task is to GUIDE them STEP BY STEP toward a specific
        answer that is DIRECTLY RELEVANT to their specific situation.
        
        IMPORTANT: Your guidance/help should ONLY be based on certain DOCUMENTS
          and NOT on your existing knowledge. NEVER answer based on your own knowledge,
          ALWAYS refer to the documents.
          However you do NOT have direct access to the docs, but you have an assistant
          named DocAgent, who DOES have access to the documents.
          
        Since you could be talking to TWO people, in order to CLARIFY who you are
        addressing, you MUST ALWAYS EXPLICITLY ADDRESS either the 
        "User" or the "DocAgent" using {AT}User or {AT}DocAgent, respectively.
        
        You must THINK like this at each step after receiving a question from the User:
        
        (I NEVER WANT TO Overwhelm DocAgent or User with TOO MANY QUESTIONS,
        so I will ALWAYS ask ONE question at a time)
        
        - I must first find out more about this topic from DocAgent, 
            let me address DocAgent to get more information.
        - I got some info from DocAgent, let me now ask the User a follow-up question
            to get ONE SPECIFIC piece of information about their situation.
        - I need to get MORE info from DocAgent, let me ask DocAgent for more info.
        - DocAgent said {NO_ANSWER}!!, Let me try asking a different way.
        - I have a bit more info, now let me ask the User a further follow-up question,
            to get ONE SPECIFIC piece of information about their situation.
        - I need more info from user, let me ask the User a follow-up question,
            to get ANOTHER SPECIFIC piece of information about their situation.
        ...[and so on]...
        - Now I have ALL the info I need from BOTH the User and DocAgent,
            so I can provide the User a DIRECTLY RELEVANT answer,
            so I will say {DONE}, followed by the answer.   
            
        IMPORTANT: When giving the User a list of choices, always show them
            a NUMBERED list of choices.          
        """,
    )
    guide_task.add_sub_task(doc_task)

    while True:
        query = Prompt.ask("[blue]How can I help?")
        if query in ["x", "q"]:
            break
        guide_task.run(query)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/doc-based-troubleshooting.py">
"""
2-agent doc-aware conversation,
different from standard question -> answer RAG flow.

User indicates some type of problem,
TroubleShooter Agent engages in conversation with User,
guiding them toward a solution.
At each step, Troubleshooter Agent can either address:
- DocAgent (who has access to docs/manuals) for info, or
- User, to ask follow-up questions about the problem

python3 examples/docqa/doc-based-troubleshooting.py

"""

import os
from typing import Optional

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import AT, DONE, NO_ANSWER

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class DocAgent(DocChatAgent):
    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        # Augment the response
        results = super().llm_response(message).content
        return self.create_llm_response(
            f"""
            Summary answer FROM DocAgent:
            
            ===
            {results}
            ===
            
            Look at the results above. These might be too much for the user to read.
            DECIDE whether you want to:
            - Ask the User a SINGLE follow-up question (could be MultipleChoice,
                where they need to select a numbered choice)
                to get more info about their situation or context, OR
            - Ask the DocAgent for more information, if you think you need more info.
            - Provide the User a FINAL answer, if you think you have enough information 
               from the User AND the Documents
               
            IMPORTANT: Do NOT simply give the User a list of options -- 
                you must HELP the user by asking them FOLLOWUP questions
                about their situation and GUIDE them to a SPECIFIC, 
                DIRECTLY RELEVANT answer. 
                You CAN give the user a MULTIPLE CHOICE question, telling them 
                to pick a number (or choice-letter) from the list.
                
            REMEMBER - NEVER ask the DocAgent or User MULTIPLE questions at a time,
                always ask ONE question at a time;
                 if asking the USER, it CAN be a MULTIPLE CHOICE question.
            """
        )


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = lm.OpenAIChatModel.GPT4o,
) -> None:
    vecdb_config = lr.vector_store.QdrantDBConfig(
        storage_path=".qdrant/doc-aware/",
        replace_collection=False,
        cloud=False,
    )

    llm_config = lm.OpenAIGPTConfig(chat_model=model)
    config = DocChatAgentConfig(
        llm=llm_config,
        vecdb=vecdb_config,
        n_query_rephrases=0,
        hypothetical_answer=False,
        assistant_mode=True,
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=100,  # aim for this many tokens per chunk
            n_neighbor_ids=5,
            overlap=20,  # overlap between chunks
            max_chunks=10_000,
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    doc_agent = DocAgent(config)
    print("[blue]Welcome to the document chatbot!")
    doc_agent.user_docs_ingest_dialog()
    print("[cyan]Enter x or q to quit, or ? for evidence")
    doc_task = Task(
        doc_agent,
        interactive=False,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
    )

    guide_agent = ChatAgent(
        ChatAgentConfig(
            name="GuideAgent",
            llm=llm_config,
            vecdb=None,
        )
    )
    # MyRecipientTool = RecipientTool.create(
    #     recipients=["DocAgent", "User"], default="User"
    # )
    # guide_agent.enable_message(MyRecipientTool)
    task_config = lr.TaskConfig(addressing_prefix=AT)
    guide_task = Task(
        guide_agent,
        interactive=False,
        config=task_config,
        system_message=f"""
        You are a TROUBLESHOOTER, who wants to help a User with their PROBLEM.
        
        Your task is to GUIDE them STEP BY STEP toward a specific
        resolution that is DIRECTLY RELEVANT to their specific problem.
        
        IMPORTANT: Your guidance/help should ONLY be based on certain DOCUMENTS
          and NOT on your existing knowledge. NEVER answer based on your own knowledge,
          ALWAYS refer to the documents.
          However you do NOT have direct access to the docs, but you have an assistant
          named DocAgent, who DOES have access to the documents.
          
        Since you could be talking to TWO people, in order to CLARIFY who you are
        addressing, you MUST ALWAYS EXPLICITLY ADDRESS either the 
        "User" or the "DocAgent" using {AT}User or {AT}DocAgent, respectively.
        
        You must THINK like this at each step after receiving a question from the User:
        
        (I NEVER WANT TO Overwhelm DocAgent or User with TOO MANY QUESTIONS,
        so I will ALWAYS ask ONE question at a time)
        
        - I must first find out more about this topic from DocAgent, 
            let me address DocAgent to get more information.
        - I got some info from DocAgent, let me now ask the User a follow-up question
            to get ONE SPECIFIC piece of information about their situation.
        - I need to get MORE info from DocAgent, let me ask DocAgent for more info.
        - DocAgent said {NO_ANSWER}!!, Let me try asking a different way.
        - I have a bit more info, now let me ask the User a further follow-up question,
            to get ONE SPECIFIC piece of information about their situation.
        - I need more info from user, let me ask the User a follow-up question,
            to get ANOTHER SPECIFIC piece of information about their situation.
        ...[and so on]...
        - Now I have ALL the info I need from BOTH the User and DocAgent,
            so I can provide the User a DIRECTLY RELEVANT answer,
            so I will say {DONE}, followed by the answer.   
            
        IMPORTANT: When giving the User a list of choices, always show them
            a NUMBERED list of choices.     
            
        I REPEAT -- NEVER use your OWN KNOWLEDGE. ALWAYS RELY ON the Documents
        from DocAgent.     
        """,
    )
    guide_task.add_sub_task(doc_task)

    while True:
        query = Prompt.ask("[blue]How can I help?")
        if query in ["x", "q"]:
            break
        guide_task.run(query)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/doc-chat-2.py">
"""
2-agent doc-chat:
WriterAgent is in charge of answering user's question.
Breaks it down into smaller questions (if needed) to send to DocAgent,
who has access to the docs via a vector-db.

python3 examples/docqa/doc-chat-2.py
"""

import os

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.agent.tools.recipient_tool import RecipientTool
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

os.environ["TOKENIZERS_PARALLELISM"] = "false"


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = lm.OpenAIChatModel.GPT4o,
) -> None:
    llm_config = lm.OpenAIGPTConfig(chat_model=model)
    config = DocChatAgentConfig(
        llm=llm_config,
        n_query_rephrases=0,
        hypothetical_answer=False,
        full_citations=False,
        assistant_mode=True,
        n_neighbor_chunks=2,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=100,  # aim for this many tokens per chunk
            n_neighbor_ids=5,
            overlap=20,  # overlap between chunks
            max_chunks=10_000,
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    embed_cfg = lr.embedding_models.OpenAIEmbeddingsConfig()

    config.vecdb = lr.vector_store.QdrantDBConfig(
        cloud=False,
        storage_path=".qdrant/doc-chat",
        embedding=embed_cfg,
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    doc_agent = DocChatAgent(config)
    print("[blue]Welcome to the document chatbot!")
    doc_agent.user_docs_ingest_dialog()
    print("[cyan]Enter x or q to quit, or ? for evidence")
    doc_task = Task(
        doc_agent,
        interactive=False,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
    )

    writer_agent = ChatAgent(
        ChatAgentConfig(
            name="WriterAgent",
            llm=llm_config,
            vecdb=None,
        )
    )
    writer_agent.enable_message(RecipientTool)
    writer_task = Task(
        writer_agent,
        name="WriterAgent",
        interactive=False,
        system_message=f"""
        You are tenacious, creative and resourceful when given a question to 
        find an answer for. You will receive questions from a user, which you will 
        try to answer ONLY based on content from certain documents (not from your 
        general knowledge). However you do NOT have access to the documents. 
        You will be assisted by DocAgent, who DOES have access to the documents.
        
        Here are the rules:
        (a) when the question is complex or has multiple parts, break it into small 
         parts and/or steps and send them to DocAgent
        (b) if DocAgent says {NO_ANSWER} or gives no answer, try asking in other ways.
        (c) Once you collect all parts of the answer, say "DONE" 
            and show me the consolidated final answer. 
        (d) DocAgent has no memory of previous dialog, so you must ensure your 
            questions are stand-alone questions that don't refer to entities mentioned 
            earlier in the dialog.
        (e) if DocAgent is unable to answer after your best efforts, you can say
            {NO_ANSWER} and move on to the next question.
        (f) answers should be based ONLY on the documents, NOT on your prior knowledge.
        (g) be direct and concise, do not waste words being polite.
        (h) if you need more info from the user, before asking DocAgent, you should 
        address questions to the "User" (not to DocAgent) to get further 
        clarifications or information. 
        (i) Always ask questions ONE BY ONE (to either User or DocAgent), NEVER 
            send Multiple questions in one message.
        (j) Use bullet-point format when presenting multiple pieces of info.
        (k) When DocAgent responds without citing a SOURCE and EXTRACT(S), you should
            send your question again to DocChat, reminding it to cite the source and
            extract(s).
        
        Start by asking the user what they want to know.
        """,
    )
    writer_task.add_sub_task(doc_task)

    while True:
        query = Prompt.ask("[blue]How can I help?")
        if query in ["x", "q"]:
            break
        writer_task.run(query)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/doc-chat-multi-llm.py">
"""
2-agent doc-chat:
WriterAgent (powered by GPT4) is in charge of answering user's question,
which can be complex.
Breaks it down into smaller questions (if needed) to send to DocAgent
(powered by a possibly weaker but cheaper LLM),
who has access to the docs via a vector-db.

You can run this with different combinations, using the -m and -mr
args to specify the LLMs for the WriterAgent and DocAgent (RAG) respectively.

See this [script](https://github.com/langroid/langroid/blob/main/examples/docqa/rag-local-simple.py)
 for examples of specifying local models.

See here for a guide on how to use Langroid with non-OpenAI LLMs (local/remote):
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import os

import typer
from rich import print

import langroid as lr
import langroid.language_models as lm
import langroid.language_models.base
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.agent.task import Task
from langroid.agent.tools.recipient_tool import RecipientTool
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name for writer agent"),
    model_rag: str = typer.Option(
        "", "--model_rag", "-mr", help="model name for RAG agent"
    ),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    llm_config_rag = OpenAIGPTConfig(
        chat_model=model_rag or model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=16_000,  # adjust based on model
        timeout=45,
    )

    config = DocChatAgentConfig(
        llm=llm_config_rag,
        n_query_rephrases=0,
        hypothetical_answer=False,
        assistant_mode=True,
        n_similar_chunks=5,
        n_relevant_chunks=5,
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=200,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            cache_type="fakeredis",
        )
    )

    doc_agent = DocChatAgent(config)
    print("[blue]Welcome to the document chatbot!")
    doc_agent.user_docs_ingest_dialog()
    print("[cyan]Enter x or q to quit, or ? for evidence")

    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[lr.Entity.LLM],
        done_if_response=[lr.Entity.LLM],
    )

    writer_agent = ChatAgent(
        ChatAgentConfig(
            name="WriterAgent",
            llm=OpenAIGPTConfig(
                chat_model=model or lm.OpenAIChatModel.GPT4o,
                chat_context_length=8192,  # adjust based on model
            ),
            vecdb=None,
        )
    )
    writer_agent.enable_message(RecipientTool)
    writer_task = Task(
        writer_agent,
        name="WriterAgent",
        system_message=f"""
        You are tenacious, creative and resourceful when given a question to 
        find an answer for. You will receive questions from a user, which you will 
        try to answer ONLY based on content from certain documents (not from your 
        general knowledge). However you do NOT have access to the documents. 
        You will be assisted by DocAgent, who DOES have access to the documents.
        
        Here are the rules:
        (a) when the question is complex or has multiple parts, break it into small 
         parts and/or steps and send them to DocAgent
        (b) if DocAgent says {NO_ANSWER} or gives no answer, try asking in other ways.
        (c) Once you collect all parts of the answer, you can say DONE and give me 
            the final answer. 
        (d) DocAgent has no memory of previous dialog, so you must ensure your 
            questions are stand-alone questions that don't refer to entities mentioned 
            earlier in the dialog.
        (e) if DocAgent is unable to answer after your best efforts, you can say
            {NO_ANSWER} and move on to the next question.
        (f) answers should be based ONLY on the documents, NOT on your prior knowledge.
        (g) be direct and concise, do not waste words being polite.
        (h) if you need more info from the user, before asking DocAgent, you should 
        address questions to the "User" (not to DocAgent) to get further 
        clarifications or information. 
        (i) Always ask questions ONE BY ONE (to either User or DocAgent), NEVER 
            send Multiple questions in one message.
        (j) Use bullet-point format when presenting multiple pieces of info.
        (k) When DocAgent responds without citing a SOURCE and EXTRACT(S), you should
            send your question again to DocChat, reminding it to cite the source and
            extract(s).
        
        
        Start by asking the user what they want to know.
        """,
    )
    writer_task.add_sub_task(doc_task)
    writer_task.run("Can you help me with some questions?")

    # show cost summary
    print("LLM usage, cost summary:")
    print(str(langroid.language_models.base.LanguageModel.usage_cost_summary()))


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/doc-chat-simple.py">
"""
Bare-bones example of using DocChatAgent to query a document.

Run like this (omit the model to use default GPT-4o):
    
    python3 examples/docqa/doc-chat-simple.py --model ollama/qwen2.5:latest
    
"""

from fire import Fire

import langroid.language_models as lm
from langroid.vector_store.chromadb import ChromaDBConfig
from langroid.agent.special.doc_chat_agent import (
    DocChatAgent,
    DocChatAgentConfig,
)


def main(model: str = ""):
    # set up the agent
    agent = DocChatAgent(
        DocChatAgentConfig(
            vecdb=ChromaDBConfig(),
            llm=lm.OpenAIGPTConfig(chat_model=model or lm.OpenAIChatModel.GPT4o),
            # several configs possible here, omitted for brevity
        ),
    )

    # ingest document(s), could be a local file/folder or URL
    # Try Borges' "Library of Babel" short story
    url = "https://xpressenglish.com/our-stories/library-of-babel/"

    agent.ingest_doc_paths([url])

    result = agent.llm_response("what is the shape of the rooms in the library?")

    assert "hexagon" in result.content.lower()

    print(result.content)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/doc-chunk-enrich.py">
"""
Demonstrating the utility of Hypothetical Questions (HQ) in the context of a
DocChatAgent.

In the following example, a DocChatAgent is created and it can be queried on its
documents both in a normal way and in a hypothetical way.

Although this is being referred to as Hypothetical Questions, it is not limited to
just questions -- it is simply a way to augment the document-chunks at ingestion time,
with keywords that increase the "semantic surface" of the chunks to improve
retrieval accuracy.

This example illustrates the benefit of HQ in a medical scenario
where each "document chunk" is simply the name of a medical test
(e.g. "cholesterol", "BUN", "PSA", etc)
and when `use_hypothetical_question` is enabled,
the chunk (i.e. test name) is augment it with keywords that add more
context, such as which organ it is related to
(e.g., "heart", "kidney", "prostate", etc).
This way, when a user asks "which tests are related to kidney health",
these augmentations ensure that the test names are retrieved more accurately.

Running the script compares the accuracy of
results of the DocChatAgent with and without HQ.

Run like this to use HQ:

python3 examples/docqa/hypothetical_questions.py

or without HQ:

python3 examples/docqa/hypothetical_questions.py --no-use-hq
"""

import typer
from rich import print
from rich.table import Table

import langroid as lr
import langroid.language_models as lm
from langroid.agent.batch import run_batch_function
from langroid.agent.special.doc_chat_agent import (
    ChunkEnrichmentAgentConfig,
    DocChatAgent,
    DocChatAgentConfig,
)
from langroid.parsing.parser import ParsingConfig
from langroid.utils.configuration import Settings
from langroid.vector_store.qdrantdb import QdrantDBConfig

app = typer.Typer()

lr.utils.logging.setup_colored_logging()

ORGAN = "kidney"


def setup_vecdb(docker: bool, reset: bool, collection: str) -> QdrantDBConfig:
    """Configure vector database."""
    return QdrantDBConfig(
        collection_name=collection, replace_collection=reset, docker=docker
    )


def run_document_chatbot(
    model: str,
    docker: bool,
    reset: bool,
    collection: str,
    use_hq: bool,
) -> None:
    """
    Main function for the document chatbot.

    Args:
        model: chat model
        docker: use docker for vector database
        reset: reset conversation memory
        collection: collection name
        use_hq: use hypothetical
    """
    llm_config = lm.OpenAIGPTConfig(chat_model=model)
    vecdb_config = setup_vecdb(docker=docker, reset=reset, collection=collection)
    enrichment_config = ChunkEnrichmentAgentConfig(
        batch_size=10,
        system_message="""
        You are an experienced clinical physician, very well-versed in
        medical tests and their names.
        You will be asked to identify WHICH ORGAN(s) Function/Health
        a test name is most closely associated with, to aid in 
        retrieving the medical test names more accurately from an embeddings db
        that contains thousands of such test names.
        The idea is to use the ORGAN NAME(S) provided by you, 
        to make the right test names easier to discover via keyword-matching
        or semantic (embedding) similarity.
         Your job is to generate up to 3 ORGAN NAMES
         MOST CLOSELY associated with the test name shown, ONE PER LINE.
         DO NOT SAY ANYTHING ELSE, and DO NOT BE OBLIGATED to provide 3 organs --
         if there is just one or two that are most relevant, that is fine.
        Examples:
          "cholesterol" -> "heart function", 
          "LDL" -> "artery health", etc,
          "PSA" -> "prostate health", 
          "TSH" -> "thyroid function", etc.                
        """,
        enrichment_prompt_fn=lambda test: f"""
        Which ORGAN(S) Function/Health is the medical test named 
        '{test}' most closely associated with?
        """,
    )

    config = DocChatAgentConfig(
        llm=llm_config,
        vecdb=vecdb_config,
        hypothetical_answer=False,
        rerank_diversity=False,
        rerank_periphery=False,
        use_reciprocal_rank_fusion=False,
        n_similar_chunks=10,
        n_relevant_chunks=10,
        parsing=ParsingConfig(
            chunk_size=120,
            overlap=15,
            min_chunk_chars=50,
        ),
        # n_neighbor_chunks=1,
        chunk_enrichment_config=enrichment_config if use_hq else None,
        relevance_extractor_config=None,
    )

    doc_agent = DocChatAgent(config=config)
    medical_tests = """
    BUN, Creatinine, GFR, ALT, AST, ALP, Albumin, Bilirubin, CBC, eGFR, PTH, 
    Uric Acid, Ammonia, Protein/Creatinine Ratio, Total Protein, LDH, SPEP, CRP, 
    ESR, Cystatin C
    """

    medical_test_list = [test.strip() for test in medical_tests.split(",")]

    # already "chunked" docs:
    docs = [lr.Document.from_string(test, is_chunk=True) for test in medical_test_list]
    # this should augment each test name with organ names that help improve retrieval
    doc_agent.ingest_docs(docs)
    if use_hq:
        print("[cyan]Test names augmented with organ names:")
        for doc in doc_agent.chunked_docs:
            print(doc.content)
            print("---")

    user_query = f"Which tests are related to {ORGAN} function?"

    _, relevant_chunks = doc_agent.get_relevant_extracts(user_query)
    relevant_chunks_str = "\n".join([chunk.content for chunk in relevant_chunks])
    print(f"relevant test names retrieved:\n{relevant_chunks_str}")
    system_msg = f"""
      You are an experienced clinical physician, well-versed in
      medical tests and their names. You are looking a set of 
      tests or readings that have been performed on a patient. 
      Based on these tests or readings, you need to determine 
      which of the tests shown are relevant to compiling a medical 
      report on the {ORGAN} function and {ORGAN} health of the 
      patient.
    """
    asst_msg = f"""
    Yes I perfectly understand! I will be diligent and discriminating, 
    and will accurately pick out which of the tests are related to
    compiling a comprehensive medical report on the {ORGAN} function and 
    {ORGAN} health. Please show me the full list of tests and/or readings
    and I PROMISE I will be able to tell you which of them are relevant to
    {ORGAN} function or {ORGAN} health.
    """
    user_msg = f"""
    Your patient had a series of tests/measurements performed,
    and below are the TEST (or measurement) NAMES that were recorded.
    For you to compile a comprehensive medical report on the {ORGAN} function and
    {ORGAN} health of the patient,
    which of these tests are typically considered related to this organ's 
    function or health? 
    
    Simply list the relevant test-names, VERBATIM exactly as they appear,
    one per line, without any explanation or elaboration.

    TESTS/MEASUREMNTS:

    {relevant_chunks_str}
    """

    retrieval_answer = doc_agent.llm.chat(
        [
            lm.LLMMessage(content=system_msg, role=lm.Role.SYSTEM),
            lm.LLMMessage(content=asst_msg, role=lm.Role.ASSISTANT),
            lm.LLMMessage(content=user_msg, role=lm.Role.USER),
        ]
    ).message
    print(f"\n\nAnswer from DocChatAgent.llm after retrieval:\n{retrieval_answer}")
    retrieval_tests = retrieval_answer.split("\n")
    retrieval_tests = [
        test.strip()
        for test in retrieval_tests
        if test.strip() and test.strip() in medical_test_list
    ]

    # compare this with directly asking the LLM about each individual test
    print(f"[blue]Directly asking the LLM whether each test is related to {ORGAN}:")
    llm = doc_agent.llm

    def llm_classify(test: str) -> str:
        return llm.chat(
            [
                lm.LLMMessage(content=system_msg, role=lm.Role.SYSTEM),
                lm.LLMMessage(content=asst_msg, role=lm.Role.ASSISTANT),
                lm.LLMMessage(
                    content=f"""
                          Is the medical test named '{test}' typically considered
                          DIRECTLY related to {ORGAN} function?,
                          simply say 'yes' or 'no'
                          """,
                    role=lm.Role.USER,
                ),
            ]
        ).message

    classifications = run_batch_function(llm_classify, medical_test_list, batch_size=5)
    direct_llm_tests = [
        test
        for test, classification in zip(medical_test_list, classifications)
        if "yes" in classification.lower()
    ]
    print("[green]Relevant tests from direct LLM query:\n")
    print("\n".join(direct_llm_tests))

    # Create a table with test comparison
    test_union = set(direct_llm_tests).union(set(retrieval_tests))

    with_str = "with" if use_hq else "without"
    table = Table(
        title=f"Test Detection Methods Comparison for {ORGAN} {with_str} Hyp Questions"
    )
    table.add_column("Test", justify="left")
    table.add_column("Direct", justify="center")
    table.add_column("Retrieval", justify="center")

    for test in sorted(test_union):
        direct = "x" if test in direct_llm_tests else ""
        retrieved = "x" if test in retrieval_tests else ""
        table.add_row(test, direct, retrieved)

    print("\n")
    print(table)

    # calc percent overlap or jacard similarity between the two sets of relevant tests
    overlap = len(
        set(direct_llm_tests).intersection(set(relevant_chunks_str.split("\n")))
    )
    union = len(test_union)
    jacard_pct = (100 * overlap / union) if union > 0 else 0
    print(
        f"[cyan]Jaccard similarity between the two sets of relevant tests: {jacard_pct:.2f}%"
    )


@app.command()
def main(
    debug: bool = typer.Option(
        False, "--debug/--no-debug", "-d", help="Enable debug mode"
    ),
    stream: bool = typer.Option(
        True, "--stream/--no-stream", "-s", help="Enable streaming output"
    ),
    cache: bool = typer.Option(True, "--cache/--no-cache", "-c", help="Enable caching"),
    model: str = typer.Option(
        lm.OpenAIChatModel.GPT4o_MINI.value, "--model", "-m", help="Chat model to use"
    ),
    collection: str = typer.Option(
        "docchat_hq", "--collection", help="Collection name for vector database"
    ),
    docker: bool = typer.Option(
        True, "--docker/--no-docker", help="Use docker for vector database"
    ),
    reset: bool = typer.Option(
        True, "--reset/--no-reset", help="Reset conversation memory"
    ),
    use_hq: bool = typer.Option(
        True, "--use-hq/--no-use-hq", help="Use hypothetical questions"
    ),
) -> None:
    """Main app function."""
    lr.utils.configuration.set_global(
        Settings(
            debug=debug,
            cache=cache,
            stream=stream,
        )
    )

    run_document_chatbot(
        model=model,
        docker=docker,
        collection=collection,
        reset=reset,
        use_hq=use_hq,
    )


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/extract-then-chat.py">
"""
3-Agent system to first extract a few pieces of info, then chat with user.

- Assistant: helps user answer questions about a Book. But first it needs to
    extract some information from a document about the Book, using Extractor.
- Extractor: generates questions about the Book document, one by one,
    then returns all info to Assistant using a tool message.
- DocAgent: answers the questions generated by Extractor, based on the Book doc.

Run like this:

python3 examples/chainlit/extract-then-chat.py

"""

import os
from typing import List

from dotenv import load_dotenv
from fire import Fire
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.mytypes import Entity
from langroid.parsing.parser import ParsingConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import DONE, NO_ANSWER, PASS, SEND_TO

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class BookInfo(BaseModel):
    title: str
    author: str
    year: int


class BookInfoTool(ToolMessage):
    request: str = "book_info"
    purpose: str = "Collect <info> about Books"

    info: List[BookInfo]

    def handle(self) -> str:
        """Exit task and pass tool to parent"""
        return DONE + " " + PASS

    @classmethod
    def examples(cls) -> List["BookInfoTool"]:
        return [
            cls(
                info=[
                    BookInfo(title="The Hobbit", author="J.R.R. Tolkien", year=1937),
                    BookInfo(
                        title="The Great Gatsby",
                        author="F. Scott Fitzgerald",
                        year=1925,
                    ),
                ]
            )
        ]


class Assistant(ChatAgent):
    def book_info(self, msg: BookInfoTool) -> str:
        # convert info  to NON-JSON so it doesn't look like a tool,
        # and insert routing so that the Assistan't LLM responds to it, not user.
        info_str = str(msg.info).replace("{", "[").replace("}", "]")
        return f"""{SEND_TO}LLM
        Below is INFO about various books, you received from the Extractor.
        Now ask the user what help they need, and respond ONLY based on this INFO.
        
        INFO: 
        {info_str} 
        """


class Extractor(ChatAgent):
    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        """Nudge LLM when it fails to use book_info correctly"""
        if self.has_tool_message_attempt(msg):
            return """
            You must use the "book_info" tool to present the info.
            You either forgot to use it, or you used it with the wrong format.
            Make sure all fields are filled out and pay attention to the 
            required types of the fields.
            """


def chat(
    model: str = "",  # or, e.g., "ollma/mistral:7b-instruct-v0.2-q8_0"
    debug: bool = False,
    no_cache: bool = False,  # whether to disablue using cached LLM responses
):
    print(
        """
        Hello! I am your book info helper. 
        First I will get info about some books
        """
    )

    load_dotenv()

    set_global(
        Settings(
            debug=debug,
            cache=not no_cache,  # disables cache lookup; set to True to use cache
        )
    )

    llm_cfg = lm.OpenAIGPTConfig(
        # or, e.g. "ollama/mistral:7b-instruct-v0.2-q8_0" but result may be brittle
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=16_000,  # adjust based on model
    )
    doc_agent = DocChatAgent(
        DocChatAgentConfig(
            llm=llm_cfg,
            n_neighbor_chunks=2,
            n_similar_chunks=3,
            n_relevant_chunks=3,
            parsing=ParsingConfig(
                chunk_size=50,
                overlap=10,
                n_neighbor_ids=4,
            ),
            vecdb=lr.vector_store.QdrantDBConfig(
                collection_name="book_info",
                replace_collection=True,
                cloud=False,
                storage_path=".qdrant/data/",
                embedding=lr.embedding_models.SentenceTransformerEmbeddingsConfig(
                    model_type="sentence-transformer",
                    model_name="BAAI/bge-large-en-v1.5",
                ),
            ),
            cross_encoder_reranking_model="",
        )
    )
    doc_agent.ingest_doc_paths(["examples/docqa/books.txt"])
    doc_task = Task(
        doc_agent,
        name="DocAgent",
        done_if_no_response=[Entity.LLM],  # done if null response from LLM
        done_if_response=[Entity.LLM],  # done if non-null response from LLM
        # Don't use system_message here since it will override doc chat agent's
        # default system message
    )

    extractor_agent = Extractor(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    extractor_agent.enable_message(BookInfoTool)

    extractor_task = Task(
        extractor_agent,
        name="Extractor",
        interactive=False,  # set to True to slow it down (hit enter to progress)
        system_message=f"""
        You are an expert at understanding JSON function/tool specifications.
        You must extract information about various books from a document,
        to finally present the info using the `book_info` tool/function,
        but you do not have access to the document. 
        I can help with your questions about the document.
        You have to ask questions in these steps:
        1. ask which books are in the document
        2. for each book, ask the various pieces of info you need.
        
        If I am unable to answer your question initially, try asking differently,
        and if I am still unable to answer after 3 tries, 
        fill in {NO_ANSWER} for that field. 
        Think step by step. 
        
        Do not explain yourself, or say any extraneous things. 
        When you receive the answer, then ask for the next field, and so on.
        """,
    )

    assistant_agent = Assistant(
        ChatAgentConfig(
            llm=llm_cfg,
            vecdb=None,
        )
    )
    assistant_agent.enable_message(lr.agent.tools.RecipientTool)
    # enable assistant to HANDLE the book_info tool but not USE it
    assistant_agent.enable_message(BookInfoTool, use=False, handle=True)
    assistant_task = Task(
        assistant_agent,
        name="Assistant",
        interactive=True,
        system_message="""
        You are a helpful librarian, answering my (the user) questions about 
        books described in a certain document, and you do NOT know which 
        books are in the document.
        
        FIRST you need to ask the "Extractor" to collect information
        about various books that are in a certain document. Address your request to the 
        Extractor using the 'recipient_message' tool/function. 
        
        Once you receive the information, you should then ask me (the user) 
        what I need help with.                
        """,
    )

    assistant_task.add_sub_task([extractor_task])
    extractor_task.add_sub_task([doc_task])

    # must use run() instead of run_async() because DocChatAgent
    # does not have an async llm_response method
    assistant_task.run()


if __name__ == "__main__":
    Fire(chat)
</file>

<file path="examples/docqa/filter-multi-doc-auto.py">
"""
Two-agent system to use to chat with multiple docs,
and use a combination of Filtering + RAG to answer questions,
where the filter is part of a query plan generated by LanceQueryPlanAgent.

Works with LanceDB vector-db.

- Main agent takes user question, generates a QueryPlan consisting of
    - filter (SQL, to use with lanceDB)
    - possibly rephrased query

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

NOTES:
(1) The app works best with GPT-4o, but results may be mixed with local LLMs.
You may have to tweak the system_message, use_message, and summarize_prompt
as indicated in comments below, to get good results.

"""

import os

import typer
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.special.doc_chat_agent import DocChatAgentConfig
from langroid.agent.special.lance_doc_chat_agent import LanceDocChatAgent
from langroid.agent.special.lance_rag.lance_rag_task import LanceRAGTaskCreator
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from pydantic import Field
from langroid.utils.configuration import Settings, set_global
from langroid.vector_store.lancedb import LanceDBConfig

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class MusicianMetadata(lr.DocMetaData):
    name: str = Field(..., description="The name of the musician.")
    birth_year: int = Field(..., description="The year the musician was born.")
    death_year: int = Field(..., description="The year the musician died.")
    type: str = Field(..., description="The type of musician, e.g. composer, musician.")
    genre: str = Field(..., description="The genre of the musician.")


class MusicianDocument(lr.Document):
    content: str
    metadata: MusicianMetadata


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=4096,  # adjust based on model
        timeout=90,
    )

    # Configs
    embed_cfg = OpenAIEmbeddingsConfig()

    # Get movies data
    COLLECTION = "chat-lance-music"
    ldb_dir = ".lancedb/data/musicians"
    ldb_cfg = LanceDBConfig(
        cloud=False,
        collection_name=COLLECTION,
        storage_path=ldb_dir,
        embedding=embed_cfg,
        replace_collection=False,
        document_class=MusicianDocument,
    )
    config = DocChatAgentConfig(
        name="MusicianBot",
        vecdb=ldb_cfg,
        n_query_rephrases=0,
        hypothetical_answer=False,
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=0,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        llm=llm_config,
        # system_message="...override default DocChatAgent system msg here",
        # user_message="...override default DocChatAgent user msg here",
        # summarize_prompt="...override default DocChatAgent summarize prompt here",
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=300,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            cache_type="fakeredis",
        )
    )

    print("[blue]Welcome to the Musician document-filtering chatbot!")

    # need a LanceDocChatAgent to use LanceRAgTaskCreator below
    agent = LanceDocChatAgent(config)

    # INGEST DOCS with META DATA
    beethoven_path = (
        "https://en.wikipedia.org/wiki/Ludwig_van_Beethoven"  # or can be local dir
    )
    mozart_path = "https://en.wikipedia.org/wiki/Wolfgang_Amadeus_Mozart"
    bach_path = "https://en.wikipedia.org/wiki/Johann_Sebastian_Bach"
    hendrix_path = "https://en.wikipedia.org/wiki/Pink_Floyd"
    prince_path = "https://en.wikipedia.org/wiki/Prince_(musician)"
    jackson_path = "https://en.wikipedia.org/wiki/Michael_Jackson"

    paths = dict(
        beethoven=beethoven_path,
        mozart=mozart_path,
        bach=bach_path,
        hendrix=hendrix_path,
        prince=prince_path,
        jackson=jackson_path,
    )

    metadata = dict(
        beethoven=MusicianMetadata(
            name="Beethoven",
            birth_year=1770,
            death_year=1827,
            type="composer",
            genre="classical",
        ),
        mozart=MusicianMetadata(
            name="Mozart",
            birth_year=1756,
            death_year=1791,
            type="composer",
            genre="classical",
        ),
        bach=MusicianMetadata(
            name="Bach",
            birth_year=1685,
            death_year=1750,
            type="composer",
            genre="classical",
        ),
        hendrix=MusicianMetadata(
            name="Hendrix",
            birth_year=1942,
            death_year=1970,
            type="musician",
            genre="rock",
        ),
        prince=MusicianMetadata(
            name="Prince",
            birth_year=1958,
            death_year=2016,
            type="musician",
            genre="rock",
        ),
        jackson=MusicianMetadata(
            name="Jackson",
            birth_year=1958,
            death_year=2009,
            type="musician",
            genre="pop",
        ),
    )

    create_collection = True
    if COLLECTION in agent.vecdb.list_collections():
        replace = Prompt.ask(
            f"Collection {COLLECTION} already exists. Replace it? (y/n)",
            choices=["y", "n"],
            default="n",
        )
        if replace == "y":
            agent.vecdb.set_collection(COLLECTION, replace=True)
        else:
            create_collection = False
    if create_collection:
        print("[blue]Ingesting docs...")
        for musician in metadata:
            agent.ingest_doc_paths(
                [paths[musician]],  # all chunks of this doc will have same metadata
                metadata[musician],
            )
        print("[blue]Done ingesting docs")

    print("[blue]Reqdy for your questions...")
    task = LanceRAGTaskCreator.new(agent, interactive=True)
    task.run("Can you help me with some questions?")


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/filter-multi-doc-manual.py">
"""
Two-agent system to use to chat with multiple docs,
and use a combination of Filtering + RAG to answer questions,
where the filter is manually set via the LanceDocChatAgentConfig.filter field.

Works with LanceDB vector-db.

- Main agent takes user question, generates a QueryPlan consisting of
    - filter (SQL, to use with lanceDB)
    - possibly rephrased query

See here for how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

NOTES:
(1) The app works best with GPT4/Turbo, but results may be mixed with local LLMs.
You may have to tweak the system_message, use_message, and summarize_prompt
as indicated in comments below, to get good results.

"""

import os

import typer
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid.agent.special.doc_chat_agent import DocChatAgentConfig
from langroid.agent.special.lance_doc_chat_agent import LanceDocChatAgent
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from pydantic import Field
from langroid.utils.configuration import Settings, set_global
from langroid.vector_store.lancedb import LanceDBConfig

app = typer.Typer()

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class MusicianMetadata(lr.DocMetaData):
    name: str = Field(..., description="The name of the musician.")
    birth_year: int = Field(..., description="The year the musician was born.")
    death_year: int = Field(..., description="The year the musician died.")
    type: str = Field(..., description="The type of musician, e.g. composer, musician.")
    genre: str = Field(..., description="The genre of the musician.")


class MusicianDocument(lr.Document):
    content: str
    metadata: MusicianMetadata


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=4096,  # adjust based on model
        timeout=90,
    )

    # Configs
    embed_cfg = OpenAIEmbeddingsConfig()

    # Get movies data
    COLLECTION = "chat-lance-music"
    ldb_dir = ".lancedb/data/musicians"
    ldb_cfg = LanceDBConfig(
        cloud=False,
        collection_name=COLLECTION,
        storage_path=ldb_dir,
        embedding=embed_cfg,
        replace_collection=False,
        document_class=MusicianDocument,
    )
    config = DocChatAgentConfig(
        name="MusicianBot",
        vecdb=ldb_cfg,
        n_query_rephrases=0,
        hypothetical_answer=False,
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=0,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        llm=llm_config,
        # system_message="...override default DocChatAgent system msg here",
        # user_message="...override default DocChatAgent user msg here",
        # summarize_prompt="...override default DocChatAgent summarize prompt here",
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=300,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            cache_type="fakeredis",
        )
    )

    print("[blue]Welcome to the Musician document-filtering chatbot!")

    # need a LanceDocChatAgent to use LanceRAgTaskCreator below
    agent = LanceDocChatAgent(config)

    # INGEST DOCS with META DATA
    beethoven_path = (
        "https://en.wikipedia.org/wiki/Ludwig_van_Beethoven"  # or can be local dir
    )
    mozart_path = "https://en.wikipedia.org/wiki/Wolfgang_Amadeus_Mozart"
    bach_path = "https://en.wikipedia.org/wiki/Johann_Sebastian_Bach"
    hendrix_path = "https://en.wikipedia.org/wiki/Pink_Floyd"
    prince_path = "https://en.wikipedia.org/wiki/Prince_(musician)"
    jackson_path = "https://en.wikipedia.org/wiki/Michael_Jackson"

    paths = dict(
        beethoven=beethoven_path,
        mozart=mozart_path,
        bach=bach_path,
        hendrix=hendrix_path,
        prince=prince_path,
        jackson=jackson_path,
    )

    metadata = dict(
        beethoven=MusicianMetadata(
            name="Beethoven",
            birth_year=1770,
            death_year=1827,
            type="composer",
            genre="classical",
        ),
        mozart=MusicianMetadata(
            name="Mozart",
            birth_year=1756,
            death_year=1791,
            type="composer",
            genre="classical",
        ),
        bach=MusicianMetadata(
            name="Bach",
            birth_year=1685,
            death_year=1750,
            type="composer",
            genre="classical",
        ),
        hendrix=MusicianMetadata(
            name="Hendrix",
            birth_year=1942,
            death_year=1970,
            type="musician",
            genre="rock",
        ),
        prince=MusicianMetadata(
            name="Prince",
            birth_year=1958,
            death_year=2016,
            type="musician",
            genre="rock",
        ),
        jackson=MusicianMetadata(
            name="Jackson",
            birth_year=1958,
            death_year=2009,
            type="musician",
            genre="pop",
        ),
    )

    create_collection = True
    if COLLECTION in agent.vecdb.list_collections():
        replace = Prompt.ask(
            f"Collection {COLLECTION} already exists. Replace it? (y/n)",
            choices=["y", "n"],
            default="n",
        )
        if replace == "y":
            agent.vecdb.set_collection(COLLECTION, replace=True)
        else:
            create_collection = False
    if create_collection:
        print("[blue]Ingesting docs...")
        for musician in metadata:
            agent.ingest_doc_paths(
                [paths[musician]],  # all chunks of this doc will have same metadata
                metadata[musician],
            )
        print("[blue]Done ingesting docs")

    musician = Prompt.ask(
        "[blue]which musician would you like to ask about?",
        choices=list(metadata.keys()),
        default="beethoven",
    )
    print(f"[blue]You chose {metadata[musician].name}")
    # this filter setting will be used by the LanceDocChatAgent
    # to restrict the docs searched from the vector-db
    config.filter = f"metadata.name = '{metadata[musician].name}'"

    print("[blue]Reqdy for your questions...")
    task = lr.Task(
        agent,
        interactive=True,
    )
    task.run("Can you help me with some questions about musicians?")


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/filter-multi-doc-query-plan.py">
"""
Single agent to chat with multiple docs, with filtering based on user query.

- user asks a query containing an implicit filter,
  e.g. "what is the birth year of Beethoven?", implying a filter on
  docs where metadata.name == "Beethoven".
- DocChatAgent answers question using RAG restricted to the filtered docs.

"""

import json
import os
from typing import Optional

from fire import Fire
from rich import print
from rich.prompt import Prompt

import langroid as lr
import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.special.doc_chat_agent import DocChatAgentConfig
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from pydantic import Field
from langroid.utils.configuration import Settings, set_global
from langroid.utils.pydantic_utils import temp_update
from langroid.vector_store.lancedb import LanceDBConfig
from langroid.vector_store.qdrantdb import QdrantDBConfig

os.environ["TOKENIZERS_PARALLELISM"] = "false"

VECDB = "qdrant"  # or "lance"


class MusicianMetadata(lr.DocMetaData):
    name: str = Field(..., description="The name of the musician.")
    birth_year: int = Field(..., description="The year the musician was born.")
    death_year: int = Field(..., description="The year the musician died.")
    type: str = Field(..., description="The type of musician, e.g. composer, musician.")
    genre: str = Field(..., description="The genre of the musician.")


class MusicianDocument(lr.Document):
    content: str
    metadata: MusicianMetadata


class QueryPlanTool(lr.ToolMessage):
    request: str = "query_plan"
    purpose: str = """
        Given a user's query, generate a query plan consisting of the <name>
        the user is asking about, (which will be used to filter the document-set)
        and a possibly modified <query> (e.g. it may not need to contain the <name>).
        """
    name: str
    query: str


class FilterDocAgent(lr.agent.special.DocChatAgent):
    def llm_response(
        self,
        message: None | str | ChatDocument = None,
    ) -> Optional[ChatDocument]:
        """Override DocChatAgent's default method,
        to call ChatAgent's llm_response, so it emits the QueryPlanTool"""
        return lr.ChatAgent.llm_response(self, message)

    def query_plan(self, msg: QueryPlanTool) -> str:
        """Handle query plan tool"""
        # Note the filter syntax depends on the type of underlying vector-db
        if VECDB == "lance":
            name_filter = f"metadata.name=='{msg.name}'"  # SQL-like syntax
        else:
            # for qdrant use this:
            name_filter_dict = dict(
                should=[dict(key="metadata.name", match=dict(value=msg.name))]
            )
            name_filter = json.dumps(name_filter_dict)
        with temp_update(self.config, {"filter": name_filter}):
            # restrict the document-set used for keyword and other non-vector
            # similarity
            self.setup_documents(filter=name_filter)
            extracts = self.get_relevant_chunks(msg.query)
        prompt = f"""
        Answer the QUESTION below based on the following EXTRACTS:
        
        EXTRACTS:
        {extracts}
        
        QUESTION: {msg.query}
        """
        response = lr.ChatAgent.llm_response(self, prompt)
        return response.content


def main(
    debug: bool = False,
    nocache: bool = False,
    model: str = "",
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        # or, other possibilities for example:
        # "litellm/bedrock/anthropic.claude-instant-v1"
        # "ollama/llama2"
        # "local/localhost:8000/v1"
        # "local/localhost:8000"
        chat_context_length=16_000,  # adjust based on model
        timeout=90,
    )

    # Configs
    embed_cfg = OpenAIEmbeddingsConfig()

    # Get movies data
    COLLECTION = "chat-filter-doc"
    # Note the filter syntax depends on the type of vecdb
    if VECDB == "lance":
        vecdb_cfg = LanceDBConfig(
            cloud=False,
            collection_name=COLLECTION,
            storage_path=".lance/data",
            embedding=embed_cfg,
            replace_collection=False,
            document_class=MusicianDocument,
        )
    else:
        vecdb_cfg = QdrantDBConfig(
            embedding=embed_cfg,
            cloud=False,
            storage_path=":memory:",  # in-memory storage
            collection_name=COLLECTION,
        )
    config = DocChatAgentConfig(
        name="MusicianBot",
        system_message="""
        You will respond to a query in 2 ways:
        
        - if you receive just a QUERY about a musician, 
            you must use the `query_plan` tool/function to generate a query plan.
        - if you receive document EXTRACTS followed by a QUESTION,
            simply answer the question based on the extracts.
            
        Start by asking the user what help they need.
        """,
        vecdb=vecdb_cfg,
        n_query_rephrases=0,
        hypothetical_answer=False,
        # set it to > 0 to retrieve a window of k chunks on either side of a match
        n_neighbor_chunks=0,
        n_similar_chunks=3,
        n_relevant_chunks=3,
        llm=llm_config,
        # system_message="...override default DocChatAgent system msg here",
        # user_message="...override default DocChatAgent user msg here",
        # summarize_prompt="...override default DocChatAgent summarize prompt here",
        parsing=ParsingConfig(  # modify as needed
            splitter=Splitter.TOKENS,
            chunk_size=300,  # aim for this many tokens per chunk
            overlap=30,  # overlap between chunks
            max_chunks=10_000,
            n_neighbor_ids=5,  # store ids of window of k chunks around each chunk.
            # aim to have at least this many chars per chunk when
            # truncating due to punctuation
            min_chunk_chars=200,
            discard_chunk_chars=5,  # discard chunks with fewer than this many chars
            # NOTE: PDF parsing is extremely challenging, each library has its own
            # strengths and weaknesses. Try one that works for your use case.
            pdf=PdfParsingConfig(
                # alternatives: "unstructured", "docling", "fitz"
                library="pymupdf4llm",
            ),
        ),
    )

    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            cache_type="fakeredis",
        )
    )

    print("[blue]Welcome to the Musician document-filtering chatbot!")

    agent = FilterDocAgent(config)
    agent.enable_message(QueryPlanTool)

    # INGEST DOCS with META DATA
    beethoven_path = (
        "https://en.wikipedia.org/wiki/Ludwig_van_Beethoven"  # or can be local dir
    )

    bach_path = "https://en.wikipedia.org/wiki/Johann_Sebastian_Bach"

    paths = dict(
        beethoven=beethoven_path,
        bach=bach_path,
    )

    metadata = dict(
        beethoven=MusicianMetadata(
            name="Beethoven",
            birth_year=1770,
            death_year=1827,
            type="composer",
            genre="classical",
        ),
        bach=MusicianMetadata(
            name="Bach",
            birth_year=1685,
            death_year=1750,
            type="composer",
            genre="classical",
        ),
    )

    create_collection = True
    if COLLECTION in agent.vecdb.list_collections():
        replace = Prompt.ask(
            f"Collection {COLLECTION} already exists. Replace it? (y/n)",
            choices=["y", "n"],
            default="n",
        )
        if replace == "y":
            agent.vecdb.set_collection(COLLECTION, replace=True)
        else:
            create_collection = False
    if create_collection:
        print("[blue]Ingesting docs...")
        for musician in metadata:
            agent.ingest_doc_paths(
                [paths[musician]],  # all chunks of this doc will have same metadata
                metadata[musician],
            )
        print("[blue]Done ingesting docs")

    print("[blue]Reqdy for your questions...")
    task = lr.Task(agent, interactive=True)
    task.run()


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/docqa/lance-rag-gh-issues.py">
"""
This example lets you ask questions about GitHub-issues for a repo.

LanceRAGTaskCreator.new(agent) takes a LanceDocChatAgent and sets up a
3-agent system with 2 additional agents:
- QueryPlanner that decides a filter, possibly rephrased query, and
  possibly also dataframe-like calculation to answer things like ("highest rated...")

- QueryPlanAnswerCritic: this looks at the QueryPlan and the answer from the RAG agent
  and suggests changes to the QueryPlan if the answer does not look satisfactory

This system combines:
- filtering using LanceDB (sql-like filtering on document fields
- semantic search using LanceDB (vector search on document content)
- Full Text Search using LanceDB (search on document content)
- Pandas-like dataframe calculations (e.g. "highest rated", "most votes", etc.)

Run like this:
    python examples/docqa/lance-rag-gh-issues.py

Optional arguments:
-nc : turn off caching (i.e. don't retrieve cached LLM responses)
-d: debug mode, to show all intermediate results
"""

import pandas as pd
import typer
from rich.prompt import Prompt

from langroid.agent.special.doc_chat_agent import DocChatAgentConfig
from langroid.agent.special.lance_doc_chat_agent import LanceDocChatAgent
from langroid.agent.special.lance_rag.lance_rag_task import LanceRAGTaskCreator
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.parsing.repo_loader import RepoLoader
from langroid.utils.configuration import Settings, set_global
from langroid.utils.system import rmdir
from langroid.vector_store.lancedb import LanceDBConfig

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    # Global settings: debug, cache
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    # Configs
    embed_cfg = OpenAIEmbeddingsConfig()

    # Get hithub issues
    ldb_dir = ".lancedb/data/gh-issues"
    rmdir(ldb_dir)
    ldb_cfg = LanceDBConfig(
        cloud=False,
        collection_name="chat-lance-gh-issues",
        storage_path=ldb_dir,
        embedding=embed_cfg,
    )

    cfg = DocChatAgentConfig(
        vecdb=ldb_cfg,
        add_fields_to_content=["state", "year", "month", "assignee", "size"],
    )
    agent = LanceDocChatAgent(cfg)
    repo = Prompt.ask(
        "Enter a GitHub repo name as owner/repo, e.g. jmorganca/ollama",
        default="jmorganca/ollama",
    )
    n_issues = Prompt.ask("How many issues to load?", default="100")

    # load github issues from a repo
    repo_loader = RepoLoader(repo)
    issues = repo_loader.get_issues(k=int(n_issues))
    issue_dicts = [iss.model_dump() for iss in issues]
    df = pd.DataFrame(issue_dicts)
    metadata_cols = []
    agent.ingest_dataframe(df, content="text", metadata=metadata_cols)

    df_description = agent.df_description

    # inform user about the df_description, in blue
    print(
        f"""
    [blue]Here's a description of the DataFrame that was ingested:
    {df_description}
    """
    )

    task = LanceRAGTaskCreator.new(agent, interactive=False)

    while True:
        question = Prompt.ask("What do you want to know? [q to quit]")
        if question == "q":
            break
        result = task.run(question)
        print(
            f"""
            Here's your answer:
            {result.content}
            """
        )


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/lance-rag-movies.py">
"""
Chat with dataset of IMDB movies.

LanceRAGTaskCreator.new(agent) takes a LanceDocChatAgent and sets up a
3-agent system with 2 additional agents:
- QueryPlanner that decides a filter, possibly rephrased query, and
  possibly also dataframe-like calculation to answer things like ("highest rated...")

- QueryPlanAnswerCritic: this looks at the QueryPlan and the answer from the RAG agent
  and suggests changes to the QueryPlan if the answer does not look satisfactory

This system combines:
- filtering using LanceDB (sql-like filtering on document fields
- semantic search using LanceDB (vector search on document content)
- Full Text Search using LanceDB (search on document content)
- Pandas-like dataframe calculations (e.g. "highest rated", "most votes", etc.)

Run like this:
    python examples/docqa/lance-rag-movies.py

Optional arguments:
-nc : turn off caching (i.e. don't retrieve cached LLM responses)
-d: debug mode, to show all intermediate results
"""

import pandas as pd
import typer
from rich import print
from rich.prompt import Prompt

from langroid.agent.special.doc_chat_agent import DocChatAgentConfig
from langroid.agent.special.lance_doc_chat_agent import LanceDocChatAgent
from langroid.agent.special.lance_rag.lance_rag_task import LanceRAGTaskCreator
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.system import rmdir
from langroid.vector_store.lancedb import LanceDBConfig

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    # Global settings: debug, cache
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            cache_type="fakeredis",
        )
    )

    # Configs
    embed_cfg = OpenAIEmbeddingsConfig()

    # Get movies data
    ldb_dir = ".lancedb/data/imdb-reviews"
    rmdir(ldb_dir)
    ldb_cfg = LanceDBConfig(
        cloud=False,
        collection_name="chat-lance-imdb",
        storage_path=ldb_dir,
        embedding=embed_cfg,
    )

    print(
        """
        [blue]Welcome to the IMDB Movies chatbot!
        This dataset has around 130,000 movie reviews, with these columns:
        
        movie, genre, runtime, certificate, rating, stars, 
        description, votes, director.
        
        To keep things speedy, we'll restrict the dataset to movies
        of a specific genre that you can choose.
        """
    )
    genre = Prompt.ask(
        "Which of these genres would you like to focus on?",
        default="Crime",
        choices=[
            "Action",
            "Adventure",
            "Biography",
            "Comedy",
            "Crime",
            "Documentary",
            "Drama",
            "Fantasy",
            "History",
            "Horror",
            "Music",
            "Musical",
            "Mystery",
            "Romance",
            "Sci-Fi",
            "Sport",
            "Thriller",
            "War",
            "Western",
        ],
    )
    cfg = DocChatAgentConfig(
        vecdb=ldb_cfg,
        add_fields_to_content=["movie", "genre", "certificate", "stars", "rating"],
        filter_fields=["genre", "certificate", "rating"],
    )
    agent = LanceDocChatAgent(cfg)

    # READ IN AND CLEAN THE DATA
    df = pd.read_csv("examples/docqa/data/movies/IMDB.csv")

    def clean_votes(value):
        """Clean the votes column"""
        # Remove commas and convert to integer, if fails return 0
        try:
            return int(value.replace(",", ""))
        except ValueError:
            return 0

    # Clean the 'votes' column
    df["votes"] = df["votes"].fillna("0").apply(clean_votes)

    # Clean the 'rating' column
    df["rating"] = df["rating"].fillna(0.0).astype(float)

    # Replace missing values in all other columns with '??'
    df.fillna("??", inplace=True)
    df["description"].replace("", "unknown", inplace=True)

    # get the rows where 'Crime' is in the genre column
    df = df[df["genre"].str.contains(genre)]

    print(
        f"""
    [blue]There are {df.shape[0]} movies in {genre} genre, hang on while I load them...
    """
    )
    # sample 1000 rows for faster testing
    df = df.sample(1000)

    # INGEST THE DataFrame into the LanceDocChatAgent
    metadata_cols = []
    agent.ingest_dataframe(df, content="description", metadata=metadata_cols)
    df_description = agent.df_description

    # inform user about the df_description, in blue
    print(
        f"""
    [blue]Here's a description of the DataFrame that was ingested:
    {df_description}
    """
    )

    task = LanceRAGTaskCreator.new(agent, interactive=False)

    while True:
        question = Prompt.ask("What do you want to know? [q to quit]")
        if question == "q":
            break
        result = task.run(question)
        print(
            f"""
            Here's your answer:
            {result.content}
            """
        )


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/langroid-lancedb-rag-movies.ipynb">
{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "authorship_tag": "ABX9TyPtHa1bpv1qlH9QN6TKgN33",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "1571259796a64a398b942576a899ef8a": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DropdownModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DropdownModel",
      "_options_labels": [
       "Action",
       "Adventure",
       "Biography",
       "Comedy",
       "Crime",
       "Documentary",
       "Drama",
       "Fantasy",
       "History",
       "Horror",
       "Music",
       "Musical",
       "Mystery",
       "Romance",
       "Sci-Fi",
       "Sport",
       "Thriller",
       "War",
       "Western"
      ],
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "DropdownView",
      "description": "Choose a genre:",
      "description_tooltip": null,
      "disabled": false,
      "index": 3,
      "layout": "IPY_MODEL_ef9065ee3d1741f594eb8dc97f9f3d07",
      "style": "IPY_MODEL_7f739e3ffaa24b2abde6d6d0b52ab003"
     }
    },
    "ef9065ee3d1741f594eb8dc97f9f3d07": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7f739e3ffaa24b2abde6d6d0b52ab003": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    }
   }
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/docqa/langroid-lancedb-rag-movies.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Retrieval-Augmented Analytics with Langroid + LanceDB\n",
    "\n",
    "\n",
    "Say you are working with a large dataset of movie ratings. Let's think about\n",
    "how to answer questions like this:\n",
    "> What is the highest rated Comedy movie about college students made after 2010?\n",
    "\n",
    "To answer this kind of question, we need:\n",
    "- filtering (on genre, year),\n",
    "- retrieval (semantic/lexical search on 'college students'),\n",
    "- computation (highest rated), and\n",
    "- LLM-based generation of the final answer.\n",
    "\n",
    "Of course, we'd like to automate the filtering and computation steps -- but how?\n",
    "\n",
    "\n",
    "We could use an LLM to generate a **Query Plan** for this --\n",
    "provided the underlying data store supports:\n",
    "- a filtering language \"known\" to LLMs (like SQL), and\n",
    "- a computation language \"known\" to LLMs (like a Pandas dataframe expression).\n",
    "\n",
    "This is where [LanceDB](https://github.com/lancedb/lancedb) (the default vector-db in Langroid) comes in:\n",
    "it's a versatile, highly performant, serverless vector-database that\n",
    "supports all of these functions within the same storage system and API:\n",
    "- Fast Full-text search (so you can do lexical search in the same store\n",
    "  where you do vector/semantic-search)\n",
    "- SQL-like metadata filtering\n",
    "- Pandas dataframe interop, so you can ingest dataframes and do pandas computations.\n",
    "**bold text**\n",
    "Leveraging Langroid's powerful Multi-Agent and tools orchestration, we built a\n",
    "3-Agent system consisting of:\n",
    "- Query Planner: Takes a user's query (like the above) and generates a Query Plan as a tool/function\n",
    "  consisting of: (a) a SQL-like filter, (b) a possibly rephrased query, and (c) an optional Pandas computation.\n",
    "- A RAG Agent (powered by LanceDB) that executes the query plan combining\n",
    "  filtering, RAG, lexical search, and optional Pandas computation.\n",
    "- A Query Plan Critic that examines the Query Plan and the RAG response, and\n",
    "  suggests improvements to the Query Planner, if any.\n",
    "\n",
    "This system can answer questions such as the above.\n",
    "You can try it out in this notebook, with a dataset of\n",
    "IMDB movie ratings.\n",
    "\n",
    "If you want to run it as a script, see here:\n",
    "https://github.com/langroid/langroid-examples/blob/main/examples/docqa/lance-rag-movies.py\n",
    "\n"
   ],
   "metadata": {
    "id": "b9fHPojfnbPy"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Install, setup, import"
   ],
   "metadata": {
    "id": "psOMvEL0Gekz"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "A8-Y_YPZutn6",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "ae2c9f85-c790-4c0f-80fc-cabd56f8a917"
   },
   "source": [
    "# Silently install, suppress all output (~2-4 mins)\n",
    "!pip install -q --upgrade langroid &> /dev/null\n",
    "!pip show langroid"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# various unfortunate things that need to be done to\n",
    "# control colab notebook behavior.\n",
    "\n",
    "# (a) output width\n",
    "\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "def set_css():\n",
    "  display(HTML('''\n",
    "  <style>\n",
    "    pre {\n",
    "        white-space: pre-wrap;\n",
    "    }\n",
    "  </style>\n",
    "  '''))\n",
    "get_ipython().events.register('pre_run_cell', set_css)\n",
    "\n",
    "# (b) logging related\n",
    "import logging\n",
    "logging.basicConfig(level=logging.ERROR)\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import logging\n",
    "for logger_name in logging.root.manager.loggerDict:\n",
    "    logger = logging.getLogger(logger_name)\n",
    "    logger.setLevel(logging.ERROR)\n",
    "\n",
    "# (c) allow async ops in colab\n",
    "!pip install nest-asyncio\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n"
   ],
   "metadata": {
    "id": "rWwH6duUzAC6",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "3947fda6-7de1-418e-eeb2-7717ef374c27"
   },
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "import pandas as pd\n",
    "from langroid.agent.special.doc_chat_agent import DocChatAgentConfig\n",
    "from langroid.agent.special.lance_doc_chat_agent import LanceDocChatAgent\n",
    "from langroid.agent.special.lance_rag.lance_rag_task import LanceRAGTaskCreator\n",
    "\n",
    "from langroid.utils.configuration import settings\n",
    "from langroid.embedding_models.models import OpenAIEmbeddingsConfig\n",
    "from langroid.vector_store.lancedb import LanceDBConfig\n",
    "settings.cache_type = \"fakeredis\"\n",
    "settings.notebook = True"
   ],
   "metadata": {
    "id": "A5N0NQwc3jX_",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "outputId": "311978cb-f35e-40db-d05b-a475006db2ae"
   },
   "execution_count": 22,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### OpenAI API Key (Needs GPT4-TURBO)"
   ],
   "metadata": {
    "id": "j-6vNfKW9J7b"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# OpenAI API Key: Enter your key in the dialog box that will show up below\n",
    "# NOTE: colab often struggles with showing this input box,\n",
    "# if so, try re-running the above cell and then this one,\n",
    "# or simply insert your API key in this cell, though it's not ideal.\n",
    "\n",
    "import os\n",
    "\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = getpass('Enter your GPT4-Turbo-capable OPENAI_API_KEY key:', stream=None)\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "id": "uvTODlZv3yyT",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "outputId": "7b9e7857-d030-4175-a2e3-551a5d807611"
   },
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Get IMDB ratings & descriptions data"
   ],
   "metadata": {
    "id": "TNsZdOjmQdgx"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# (1) Get the movies dataset\n",
    "\n",
    "import requests\n",
    "file_url = \"https://raw.githubusercontent.com/langroid/langroid-examples/main/examples/docqa/data/movies/IMDB.csv\"\n",
    "response = requests.get(file_url)\n",
    "with open('movies.csv', 'wb') as file:\n",
    "    file.write(response.content)\n",
    "\n"
   ],
   "metadata": {
    "id": "fegAio3kpgoo",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "outputId": "140daf33-e39d-403e-c5f1-f8225fe2ad10"
   },
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    " print(\n",
    "        \"\"\"\n",
    "        Welcome to the IMDB Movies chatbot!\n",
    "        This dataset has around 130,000 movie reviews, with these columns:\n",
    "\n",
    "        movie, genre, runtime, certificate, rating, stars,\n",
    "        description, votes, director.\n",
    "\n",
    "        To keep things speedy, we'll restrict the dataset to movies\n",
    "        of a specific genre that you can choose.\n",
    "        \"\"\"\n",
    "    )"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 191
    },
    "id": "J_Mv32fpOhgH",
    "outputId": "305849bf-0eb4-45c3-bce7-6839462c793c"
   },
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "from ipywidgets import Dropdown\n",
    "genres = [\n",
    "          \"Action\",\n",
    "          \"Adventure\",\n",
    "          \"Biography\",\n",
    "          \"Comedy\",\n",
    "          \"Crime\",\n",
    "          \"Documentary\",\n",
    "          \"Drama\",\n",
    "          \"Fantasy\",\n",
    "          \"History\",\n",
    "          \"Horror\",\n",
    "          \"Music\",\n",
    "          \"Musical\",\n",
    "          \"Mystery\",\n",
    "          \"Romance\",\n",
    "          \"Sci-Fi\",\n",
    "          \"Sport\",\n",
    "          \"Thriller\",\n",
    "          \"War\",\n",
    "          \"Western\",\n",
    "      ]\n",
    "dropdown = Dropdown(options=genres, value=genres[0], description=\"Choose a genre:\", disabled=False)\n",
    "display(dropdown)\n",
    "genre = dropdown.value"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "1571259796a64a398b942576a899ef8a",
      "ef9065ee3d1741f594eb8dc97f9f3d07",
      "7f739e3ffaa24b2abde6d6d0b52ab003"
     ]
    },
    "id": "eRNHcFHALi67",
    "outputId": "8e0936cd-7522-438c-afda-39717affa513"
   },
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# READ IN AND CLEAN THE DATA\n",
    "import pandas as pd\n",
    "df = pd.read_csv(\"movies.csv\")\n",
    "\n",
    "def clean_votes(value):\n",
    "    \"\"\"Clean the votes column\"\"\"\n",
    "    # Remove commas and convert to integer, if fails return 0\n",
    "    try:\n",
    "        return int(value.replace(\",\", \"\"))\n",
    "    except ValueError:\n",
    "        return 0\n",
    "\n",
    "# Clean the 'votes' column\n",
    "df[\"votes\"] = df[\"votes\"].fillna(\"0\").apply(clean_votes)\n",
    "\n",
    "# Clean the 'rating' column\n",
    "df[\"rating\"] = df[\"rating\"].fillna(0.0).astype(float)\n",
    "\n",
    "# Replace missing values in all other columns with '??'\n",
    "df.fillna(\"??\", inplace=True)\n",
    "df[\"description\"].replace(\"\", \"unknown\", inplace=True)\n",
    "\n",
    "# get the rows with selected genre\n",
    "df = df[df[\"genre\"].str.contains(genre)]\n",
    "\n",
    "print(\n",
    "    f\"\"\"\n",
    "[blue]There are {df.shape[0]} movies in {genre} genre, hang on while I load them...\n",
    "\"\"\"\n",
    ")\n",
    "# sample 1000 rows for faster testing\n",
    "df = df.sample(1000)"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 69
    },
    "id": "oBpuyfowOE7a",
    "outputId": "273e7074-2588-482d-b006-d9796f270d7c"
   },
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Set up LanceDB Vector-DB and LanceDocChatAgent"
   ],
   "metadata": {
    "id": "5rUPu_WVQprG"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Config LanceDB vector database\n",
    "import shutil\n",
    "db_dir = \".lancedb/data\"\n",
    "shutil.rmtree(db_dir)\n",
    "ldb_cfg = LanceDBConfig(\n",
    "    collection_name=\"chat-lance-imdb\",\n",
    "    replace_collection=True,\n",
    "    storage_path=db_dir,\n",
    "    embedding=OpenAIEmbeddingsConfig()\n",
    ")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "hPDrNYNzLJMt",
    "outputId": "6cd7f681-cedc-450e-baa8-5aab3fdd1389"
   },
   "execution_count": 17,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# configure, create LanceDocChatAgent\n",
    "cfg = DocChatAgentConfig(\n",
    "        vecdb=ldb_cfg,\n",
    "        show_stats=False,\n",
    "        add_fields_to_content=[\"movie\", \"genre\", \"certificate\", \"stars\", \"rating\"],\n",
    "        filter_fields=[\"genre\", \"certificate\", \"rating\"],\n",
    "    )\n",
    "agent = LanceDocChatAgent(cfg)\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "id": "SXpnQCV4MF4T",
    "outputId": "82040764-a60c-418d-cfaa-7a1711a9ac14"
   },
   "execution_count": 18,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Ingest data into LanceDocChatAgent"
   ],
   "metadata": {
    "id": "gArmf8GhQxC-"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Ingest the data into LanceDocChatAgent\n",
    "agent.ingest_dataframe(df, content=\"description\", metadata=[])\n",
    "df_description = agent.df_description\n",
    "\n",
    "# inform user about the df_description, in blue\n",
    "print(\n",
    "    f\"\"\"\n",
    "Here's a description of the DataFrame that was ingested:\n",
    "{df_description}\n",
    "\"\"\"\n",
    ")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 416
    },
    "id": "JttFhOw-MSX9",
    "outputId": "a1279c79-3574-4c69-ec47-15cb1aa4e57f"
   },
   "execution_count": 19,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Create, run a 3-agent system to handle user queries\n"
   ],
   "metadata": {
    "id": "BZcvWNXDO4gt"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "task = LanceRAGTaskCreator.new(agent, interactive=True)\n",
    "\n",
    "task.run(\"Can you help with some questions about these movies?\")"
   ],
   "metadata": {
    "id": "nVrqsGNFOyG4"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "id": "xOTmfjXjPBn4"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}
</file>

<file path="examples/docqa/lease.txt">
EX-10 2 elmonteleaseforfiling.htm MATERIAL CONTRACT
COMMERCIAL LEASE AGREEMENT



THIS LEASE AGREEMENT is made and entered into on December 1, 2013, by and between Temple CB, LLC, whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as "Landlord"), and Okra Energy, Inc., whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as "Tenant").



ARTICLE I - GRANT OF LEASE



Landlord, in consideration of the rents to be paid and the covenants and agreements to be performed and observed by the Tenant, does hereby lease to the Tenant and the Tenant does hereby lease and take from the Landlord the property described in Exhibit "A" attached hereto and by reference made a part hereof (the "Leased Premises"), together with, as part of the parcel, all improvements located thereon.



ARTICLE II - LEASE TERM



Section l.  Term of Lease.  The term of this Lease shall begin on the Commencement Date, as defined in Section 2 of this Article II, and shall terminate on May 31, 2020 ("the Termination Date"); provided, however, that at the option of Tenant, Tenant may renew this Lease for five additional successive one- year terms at a Monthly Rent of $100,000 per month, provided that notice of such renewal is given in writing no less than 120 days prior to the Termination Date or the expiration of any one-year renewal term. Tenant may at any time cancel this Lease and terminate all of its obligations hereunder by the payment of $300,000, plus all other amounts then due under this Lease.



Section 2.  Commencement Date. The "Commencement Date" shall mean  December 1, 2013.



ARTICLE III - EXTENSIONS



The parties hereto may elect to extend this Agreement upon such terms and conditions as may be agreed upon in writing and signed by the parties at the time of any such extension.



ARTICLE IV - DETERMINATION OF RENT



Section 1. Monthly Rent: The Tenant agrees to pay the Landlord and the Landlord agrees to accept, during the term hereof, at such place as the Landlord shall from time to time direct by notice to the Tenant, monthly rent of $100,000.


Section 2.  Late Fee.  A late fee in the amount of 5% of the Monthly Rent shall be assessed if payment is not postmarked or received by Landlord on or before the tenth day of each month.



ARTICLE V - SECURITY DEPOSIT



The Tenant has deposited with the Landlord the sum of Twenty Thousand Dollars ($20,000.00) as security for the full and faithful performance by the Tenant of all the terms of this lease required to be performed by the Tenant. Such sum shall be returned to the Tenant after the expiration of this lease, provided the Tenant has fully and faithfully carried out all of its terms. In the event of a bona fide sale of the property of which the leased premises are a part, the Landlord shall have the right to transfer the security to the purchaser to be held under the terms of this lease, and the Landlord shall be released from all liability for the return of such security to the Tenant.



ARTICLE VI - TAXES



Section l.  Personal Property Taxes.  The Tenant shall be liable for all taxes levied against any leasehold interest of the Tenant or personal property and trade fixtures owned or placed by the Tenant in the Leased Premises.



Section 2.  Real Estate Taxes.  During the continuance of this lease Landlord shall deliver to Tenant a copy of any real estate taxes and assessments against the Leased Property. From and after the Commencement Date, the Tenant shall pay to Landlord not later than twenty-one (21) days after the day on which the same may become initially due, all real estate taxes and assessments applicable to the Leased Premises, together with any interest and penalties lawfully imposed thereon as a result of Tenant's late payment thereof, which shall be levied upon the Leased Premises during the term of this Lease.



Section 3.  Contest of Taxes.  The Tenant, at its own cost and expense, may, if it shall in good faith so desire, contest by appropriate proceedings the amount of any personal or real property tax. The Tenant may, if it shall so desire, endeavor at any time or times, by appropriate proceedings, to obtain a reduction in the assessed valuation of the Leased Premises for tax purposes. In any such event, if the Landlord agrees, at the request of the Tenant, to join with the Tenant at Tenant's expense in said proceedings and the Landlord agrees to sign and deliver such papers and instruments as may be necessary to prosecute such proceedings, the Tenant shall have the right to contest the amount of any such tax and the Tenant shall have the right to withhold payment of any such tax, if the statute under which the Tenant is contesting such tax so permits.



Section 4.  Payment of Ordinary Assessments.  The Tenant shall pay all assessments, ordinary and extraordinary, attributable to or against the Leased Premises not later than twenty-one (21) days after the day on which the same became initially due. The Tenant may take the benefit of any law allowing assessments to be paid in installments and in such event the Tenant shall only be liable for such installments of assessments due during the term hereof.
</file>

<file path="examples/docqa/oai-multi-extract.py">
"""
Two-agent chat with Retrieval-augmented LLM + function-call/tool.
ExtractorAgent (has no access to docs) is tasked with extracting structured
information from a commercial lease document, and must present the terms in
a specific nested JSON format.
DocAgent (has access to the lease) helps answer questions about the lease.
Repeat: WriterAgent --Question--> DocAgent --> Answer

Example:
python3 examples/docqa/chat_multi_extract.py

Use -f option to use OpenAI function calling API instead of Langroid tool.
"""

import json
import os

import typer
from rich import print

import langroid as lr
from langroid.agent.openai_assistant import (
    AssistantTool,
    OpenAIAssistant,
    OpenAIAssistantConfig,
)
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from pydantic import BaseModel
from langroid.utils.constants import DONE, NO_ANSWER
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()
os.environ["TOKENIZERS_PARALLELISM"] = "false"


class LeasePeriod(BaseModel):
    start_date: str
    end_date: str


class LeaseFinancials(BaseModel):
    monthly_rent: str
    deposit: str


class Lease(BaseModel):
    """
    Various lease terms.
    Nested fields to make this more interesting/realistic
    """

    period: LeasePeriod
    financials: LeaseFinancials
    address: str


class LeaseMessage(ToolMessage):
    """Tool/function to use to present details about a commercial lease"""

    request: str = "lease_info"
    purpose: str = "Collect information about a Commercial Lease."
    terms: Lease

    def handle(self):
        """Handle this tool-message when the LLM emits it.
        Under the hood, this method is transplated into the OpenAIAssistant class
        as a method with name `lease_info`.
        """
        print(f"DONE! Successfully extracted Lease Info:" f"{self.terms}")
        return DONE + " " + json.dumps(self.terms.model_dump())


@app.command()
def chat() -> None:
    retriever_cfg = OpenAIAssistantConfig(
        name="LeaseRetriever",
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
        system_message="Answer questions based on the documents provided.",
    )

    retriever_agent = OpenAIAssistant(retriever_cfg)
    retriever_agent.add_assistant_tools([AssistantTool(type="retrieval")])
    retriever_agent.add_assistant_files(["examples/docqa/lease.txt"])

    retriever_task = Task(
        retriever_agent,
        interactive=False,
        done_if_response=[lr.Entity.LLM],
        done_if_no_response=[lr.Entity.LLM],
    )

    extractor_cfg = OpenAIAssistantConfig(
        name="LeaseExtractor",
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
        system_message=f"""
        You have to collect information about a Commercial Lease from a 
        lease contract which you don't have access to. You need to ask
        questions ONE BY ONE to get this information. 
        Once you have all the REQUIRED fields, 
        you have to present it to me using the `lease_info` 
        function/tool (fill in {NO_ANSWER} for slots that you are unable to fill).
        """,
    )
    extractor_agent = OpenAIAssistant(extractor_cfg)
    extractor_agent.enable_message(LeaseMessage, include_defaults=False)

    extractor_task = Task(extractor_agent, interactive=False)
    extractor_task.add_sub_task(retriever_task)
    extractor_task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/oai-retrieval-2.py">
"""
Use TWO OpenAI Assistants in Langroid's Multi-Agent mode to answer questions:
 - Planner Agent: takes user question, plans, decides how to ask the Retrieval Agent
 - Retrieval Agent: takes the question from the Master Agent, answers based on docs

Run like this:
python3 examples/docqa/oai-retrieval-2.py

"""

import os
import tempfile

import typer
from rich import print
from rich.prompt import Prompt

from langroid.agent.openai_assistant import (
    AssistantTool,
    OpenAIAssistant,
    OpenAIAssistantConfig,
)
from langroid.agent.task import Task
from langroid.agent.tools.recipient_tool import RecipientTool
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.url_loader import URLLoader
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()
os.environ["TOKENIZERS_PARALLELISM"] = "false"


@app.command()
def chat() -> None:
    reuse = (
        Prompt.ask(
            "Reuse existing assistant, threads if available? (y/n)",
            default="y",
        )
        == "y"
    )

    planner_cfg = OpenAIAssistantConfig(
        name="Planner",
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
        use_cached_thread=reuse,
        use_cached_assistant=reuse,
        system_message="""
        You will receive questions from the user about some docs, 
        but you don't have access to them, but you have a Retriever to help you, since
        they have access to the docs. For each question I send you, decide how you want 
        to ask the Retriever: you can rephrase, decompose or simplify the question and 
        send it to the retriever. Once you think you have the info I need, then send 
        me (the User) a message with your consolidated answer, starting with "ANSWER:"    
        
        Start by greeting the user and asking what they want to know.     
        """,
    )
    planner_agent = OpenAIAssistant(planner_cfg)
    planner_agent.enable_message(RecipientTool)

    retriever_cfg = OpenAIAssistantConfig(
        name="Retriever",
        use_cached_thread=reuse,
        use_cached_assistant=reuse,
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
        system_message="Answer questions based on the documents provided.",
    )

    retriever_agent = OpenAIAssistant(retriever_cfg)

    print("[blue]Welcome to the retrieval chatbot!")
    path = Prompt.ask("Enter a URL or file path")
    # if path is a url, use UrlLoader to get text as a document
    if path.startswith("http"):
        text = URLLoader([path]).load()[0].content
        # save text to a temp file
        with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
            f.write(text)
            f.close()
            # get the filename
            path = f.name
    retriever_agent.add_assistant_tools([AssistantTool(type="retrieval")])
    if path:  # path may be empty if continuing from previous session
        retriever_agent.add_assistant_files([path])

    print("[cyan]Enter x or q to quit")

    planner_task = Task(planner_agent, interactive=True)

    retriever_task = Task(
        retriever_agent,
        interactive=False,
        done_if_response=[Entity.LLM],
        done_if_no_response=[Entity.LLM],
    )
    planner_task.add_sub_task(retriever_task)
    planner_task.run("")


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/oai-retrieval-assistant.py">
"""
Use OpenAI Assistant with Retrieval tool + file to answer questions.

Run like this:
python3 examples/docqa/oai-retrieval-assistant.py

"""

import os
import tempfile

import typer
from rich import print
from rich.prompt import Prompt

from langroid.agent.openai_assistant import (
    AssistantTool,
    OpenAIAssistant,
    OpenAIAssistantConfig,
)
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.parsing.url_loader import URLLoader
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()
os.environ["TOKENIZERS_PARALLELISM"] = "false"


@app.command()
def chat() -> None:
    reuse = (
        Prompt.ask(
            "Reuse existing assistant, threads if available? (y/n)",
            default="y",
        )
        == "y"
    )

    cfg = OpenAIAssistantConfig(
        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
        use_cached_assistant=reuse,
        use_cached_thread=reuse,
        system_message="Answer questions based on the provided document.",
    )
    agent = OpenAIAssistant(cfg)

    print("[blue]Welcome to the retrieval chatbot!")
    path = Prompt.ask("Enter a URL or file path")
    # if path is a url, use UrlLoader to get text as a document
    if path.startswith("http"):
        text = URLLoader([path]).load()[0].content
        # save text to a temp file
        with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
            f.write(text)
            f.close()
            # get the filename
            path = f.name
    agent.add_assistant_tools([AssistantTool(type="retrieval")])

    if path:  # may be empty if continuing from previous session
        agent.add_assistant_files([path])

    print("[cyan]Enter x or q to quit")

    task = Task(agent)

    task.run("Please help me with questions about the document I provided")


if __name__ == "__main__":
    app()
</file>

<file path="examples/docqa/rag-concurrent.py">
"""
Concurrent RAG example using DocChatAgent with custom asyncio harness

This example demonstrates running multiple DocChat queries concurrently
with detailed live logging that shows every task starting and finishing
in real time (no waiting for gather() to return), making concurrency
easy to verify at a glance.

IMPORTANT: The --sequential flag runs tasks in a TRUE sequential loop
(not asyncio's sequential mode), providing a baseline for comparison.

Usage:

# Run concurrently with asyncio (default)
python3 examples/docqa/rag-concurrent.py

# Run in TRUE sequential mode (simple loop) for baseline comparison
python3 examples/docqa/rag-concurrent.py --sequential

# With specific model
python3 examples/docqa/rag-concurrent.py -m ollama/mistral:7b-instruct-v0.2-q8_0

# Use local SentenceTransformer embeddings with Docker Qdrant on localhost:6333
python3 examples/docqa/rag-concurrent.py --local-embeddings

# Turn on cross-encoder reranking (auto-picks CUDA/MPS/CPU; override with device flag)
python3 examples/docqa/rag-concurrent.py --cross-encoder
python3 examples/docqa/rag-concurrent.py --cross-encoder --cross-encoder-device=mps

# Compare both modes to measure concurrency speedup
python3 examples/docqa/rag-concurrent.py --sequential  # Baseline
python3 examples/docqa/rag-concurrent.py  # Should be faster if truly concurrent

# Use Langroid's built-in run_batch_tasks harness instead of the custom one
python3 examples/docqa/rag-concurrent.py --use-builtin-batch

# Show only concurrency logs (suppress long answers) and filter to START/WORKER lines
python3 examples/docqa/rag-concurrent.py --num-questions=3 --log-only \\
  | rg "Q[0-9]{2} (START|WORKER|COMPLETE)"

The logs show:
- Timestamps (HH:MM:SS.mmm) for each task start/complete
- Thread IDs to verify parallel execution
- Question numbers for tracking

Expected patterns:
- SEQUENTIAL: START->COMPLETE->START->COMPLETE (one at a time)
- CONCURRENT: Multiple STARTs with close timestamps before any COMPLETEs

If concurrent mode shows START->COMPLETE pattern, there's a bottleneck
(e.g., shared vecdb client causing serialization).

See here for more on how to set up a local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import asyncio
import os
import threading
import time
from contextvars import ContextVar
from datetime import datetime
from typing import Dict

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.batch import run_batch_task_gen
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Thread-safe logging with timestamps
log_lock = threading.Lock()
CURRENT_QUESTION: ContextVar[int | None] = ContextVar("CURRENT_QUESTION", default=None)
EVENT_HISTORY: list[str] = []
QUESTION_TO_INDEX: Dict[str, int] = {}


def log_event(event_type: str, question_num: int, message: str = ""):
    """Thread-safe logging with precise timestamps"""
    timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
    thread_id = threading.get_ident() % 10000  # Short thread ID
    line = (
        f"[{timestamp}] [{thread_id:04d}] "
        f"Q{question_num:02d} {event_type:12s} {message}"
    )
    EVENT_HISTORY.append(line)
    with log_lock:
        print(line)


# 10 questions about Borges' "The Library of Babel"
ALL_QUESTIONS = [
    "What is the structure of the Library described in the story?",
    "What do the books in the Library contain?",
    "What is the significance of the hexagonal galleries?",
    "How many books are estimated to exist in the Library?",
    "What is the narrator's theory about the origin of the Library?",
    "How does the story describe the contents of most books?",
    "What happens to librarians who search for meaningful books?",
    "What is the emotional impact of the infinite Library on the librarians?",
    "What philosophical themes does the story explore?",
    "What is the relationship between infinity and meaning in the story?",
]


class LoggingDocChatAgent(DocChatAgent):
    """DocChatAgent that reports worker-thread execution for visibility."""

    def answer_from_docs(self, query: str):
        q_num = CURRENT_QUESTION.get()
        if q_num is None:
            q_num = QUESTION_TO_INDEX.get(query)
        if q_num is not None:
            log_event(
                "WORKER_START", q_num, f"Vec/LLM on T{threading.get_ident()%10000:04d}"
            )
            start = time.time()
        result = super().answer_from_docs(query)
        if q_num is not None:
            elapsed = time.time() - start
            log_event(
                "WORKER_DONE",
                q_num,
                f"{elapsed:.2f}s on T{threading.get_ident()%10000:04d}",
            )
        return result


def app(
    m: str = "",
    sequential: bool = False,
    num_questions: int = 10,
    log_only: bool = False,
    use_builtin_batch: bool = False,
    local_embeddings: bool = False,
    cross_encoder: bool = False,
    cross_encoder_device: str = "",
):
    """
    Run DocChat queries on Library of Babel story.

    Args:
        m: Model name (default: GPT-4o)
        sequential: If True, run truly sequentially (simple loop);
                   if False, run with asyncio concurrency (default: False)
        num_questions: Number of questions to run (max 10)
        log_only: Suppress verbose answers and print a concise log summary
        use_builtin_batch: Use Langroid's run_batch_tasks instead of the custom harness
        cross_encoder: Enable reranking via cross encoder (auto-picks CUDA/MPS/CPU)
        cross_encoder_device: Optional explicit device override (e.g. "cuda", "mps")
    """
    num_questions = max(1, min(num_questions, len(ALL_QUESTIONS)))
    questions = ALL_QUESTIONS[:num_questions]
    QUESTION_TO_INDEX.clear()
    QUESTION_TO_INDEX.update({q: i + 1 for i, q in enumerate(questions)})
    EVENT_HISTORY.clear()
    mode = "TRULY SEQUENTIAL (simple loop)" if sequential else "CONCURRENT (asyncio)"
    print(f"\n{'='*80}")
    print(f"Running in {mode} mode")
    print(f"{'='*80}\n")

    # Create the llm config object
    llm_config = lm.OpenAIGPTConfig(
        chat_model=m or lm.OpenAIChatModel.GPT4o,
        chat_context_length=32_000,
        max_output_tokens=300,
        temperature=0.2,
        stream=False,  # Disable streaming for batch processing
        timeout=45,
    )

    # Configure DocChatAgent with Library of Babel story
    vecdb_config = None
    if local_embeddings:
        try:
            from langroid.embedding_models.models import (
                SentenceTransformerEmbeddingsConfig,
            )
            from langroid.vector_store.qdrantdb import QdrantDBConfig
        except ImportError as exc:
            raise RuntimeError(
                "SentenceTransformer embeddings require the hf-embeddings extras"
            ) from exc

        os.environ.setdefault("QDRANT_API_URL", "http://localhost:6333")
        os.environ.setdefault("QDRANT_API_KEY", "local-dev-key")

        sentence_cfg = SentenceTransformerEmbeddingsConfig(
            model_type="sentence-transformer",
            model_name="sentence-transformers/all-MiniLM-L6-v2",
        )
        vecdb_config = QdrantDBConfig(
            cloud=True,
            collection_name="doc-chat-local-embeddings",
            replace_collection=True,
            embedding=sentence_cfg,
        )

    config_kwargs = dict(
        name="RagAgent",
        llm=llm_config,
        relevance_extractor_config=None,
    )
    if vecdb_config is not None:
        config_kwargs["vecdb"] = vecdb_config

    if cross_encoder:
        config_kwargs.update(
            dict(
                cross_encoder_reranking_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
                cross_encoder_device=cross_encoder_device or None,
            )
        )

    config = DocChatAgentConfig(**config_kwargs)

    # Create agent and ingest the document
    agent = LoggingDocChatAgent(config)
    url = "https://xpressenglish.com/our-stories/library-of-babel/"
    print(f"\nIngesting document: {url}")
    agent.ingest_doc_paths([url])
    print("Document ingested successfully.\n")
    if local_embeddings and agent.vecdb is not None:
        agent.vecdb.config.replace_collection = False

    # Create a single task that will be cloned for each question
    print(f"Creating task for concurrent execution of {len(questions)} queries...\n")

    task = lr.Task(
        agent,
        interactive=False,
        single_round=True,
    )

    # Run tasks and measure time
    print("\n" + "=" * 80)
    print("EXECUTION LOG (with timestamps and thread IDs)")
    print("=" * 80 + "\n")
    start_time = time.time()

    if sequential:
        # TRUE SEQUENTIAL: Simple loop, no async
        results = []
        for i, question in enumerate(questions, 1):
            log_event(
                "START", i, question[:50] + "..." if len(question) > 50 else question
            )  # noqa: E501
            token = CURRENT_QUESTION.set(i)
            try:
                result = task.run(question, turns=1)
            finally:
                CURRENT_QUESTION.reset(token)
            log_event(
                "COMPLETE",
                i,
                f"Got response ({len(str(result.content if result else ''))} chars)",
            )  # noqa: E501
            final = (
                result.content
                if result and hasattr(result, "content")
                else str(result) if result else ""
            )  # noqa: E501
            results.append(final)
    else:
        if use_builtin_batch:

            def input_map(question: str) -> str:
                q_num = QUESTION_TO_INDEX[question]
                log_event(
                    "START",
                    q_num,
                    question[:50] + "..." if len(question) > 50 else question,
                )
                return question

            # run_batch_task_gen allows handle_exceptions to crash on errors
            def gen_task(i: int) -> lr.Task:
                return task.clone(i)

            raw_results_gen = run_batch_task_gen(
                gen_task=gen_task,
                items=questions,
                input_map=input_map,
                sequential=False,
                turns=1,
                handle_exceptions=False,  # Crash on errors to see what's failing
            )
            raw_results = list(raw_results_gen)
            results = []
            for i, result in enumerate(raw_results, 1):
                if result is None:
                    length = 0
                    results.append("")
                elif hasattr(result, "content"):
                    length = len(result.content)
                    results.append(result)
                else:
                    text = str(result)
                    length = len(text)
                    results.append(text)
                log_event("COMPLETE", i, f"Got response ({length} chars)")
        else:
            # CONCURRENT: Custom asyncio runner using task clones and as_completed

            async def run_question(clone_idx: int, question: str, base_task: lr.Task):
                """Launch a clone of the base task and report progress live."""
                q_num = clone_idx + 1
                log_event(
                    "START",
                    q_num,
                    question[:50] + "..." if len(question) > 50 else question,
                )
                token = CURRENT_QUESTION.set(q_num)
                task_clone = base_task.clone(clone_idx)
                try:
                    result = await task_clone.run_async(question, turns=1)
                finally:
                    CURRENT_QUESTION.reset(token)
                if result is None:
                    length = 0
                elif hasattr(result, "content"):
                    length = len(result.content)
                else:
                    length = len(str(result))
                log_event("COMPLETE", q_num, f"Got response ({length} chars)")
                return q_num, result

            async def run_all_concurrent():
                coros = [
                    run_question(idx, question, task)
                    for idx, question in enumerate(questions)
                ]
                results_ordered = [None] * len(questions)
                for coro in asyncio.as_completed(coros):
                    q_num, result = await coro
                    results_ordered[q_num - 1] = result
                return results_ordered

            results = asyncio.run(run_all_concurrent())

    elapsed_time = time.time() - start_time
    print(f"\n{'='*80}")
    print(f"Completed {len(questions)} queries in {elapsed_time:.2f} seconds")
    print(f"Average time per query: {elapsed_time/len(questions):.2f} seconds")
    print(f"{'='*80}\n")

    if log_only:
        print("\nLOG SUMMARY (captured START/WORKER/COMPLETE events)")
        print("-" * 80)
        for line in EVENT_HISTORY:
            print(line)
        print("-" * 80)
    else:
        print("\nINTERPRETING THE LOGS:")
        print("-" * 80)
        if sequential:
            print("SEQUENTIAL MODE: Tasks run one at a time in a simple loop")
            print("You should see: START->COMPLETE->START->COMPLETE pattern")
            print("This is the baseline for comparison.")
        else:
            print("CONCURRENT MODE: Tasks should run in parallel with asyncio")
            print(
                "Expected: Multiple 'START' events with close timestamps "
                "BEFORE any 'COMPLETE'"
            )
            print("If you see START->COMPLETE->START->COMPLETE instead,")
            print(
                "then there's a bottleneck preventing concurrency (e.g., shared vecdb)"
            )
            print("\nThread IDs: Different IDs = parallel execution")
            print("Timestamps: Overlapping windows = true concurrency")
        print("-" * 80 + "\n")

        # Display results
        for i, (question, result) in enumerate(zip(questions, results), 1):
            print(f"\n{'='*80}")
            print(f"Q{i}: {question}")
            print(f"{'-'*80}")
            if result is not None:
                answer = result.content if hasattr(result, "content") else str(result)
            else:
                answer = "No response"
            print(f"A{i}: {answer}")
            print(f"{'='*80}")

    QUESTION_TO_INDEX.clear()
    return results


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/docqa/rag-local-simple.py">
"""
RAG example using a local LLM, with ollama

Run like this --

python3 examples/docqa/rag-local-simple.py -m <model_name>

For example, you can get good results using:
```
ollama run mistral:7b-instruct-v0.2-q8_0

python3 examples/docqa/rag-local-simple.py -m ollama/mistral:7b-instruct-v0.2-q8_0


See here for more on how to set up a local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import os

import fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig

os.environ["TOKENIZERS_PARALLELISM"] = "false"


def app(m="ollama/mistral:7b-instruct-v0.2-q8_0"):
    # Create the llm config object.
    llm_config = lm.OpenAIGPTConfig(
        # if you comment out `chat_model`, it will default to OpenAI GPT4-turbo
        # chat_model="ollama/mistral:7b-instruct-v0.2-q4_K_M",
        chat_model=m or lm.OpenAIChatModel.GPT4o,
        chat_context_length=32_000,  # set this based on model
        max_output_tokens=100,
        temperature=0.2,
        stream=True,
        timeout=45,
    )

    # Recommended: First test if basic chat works with this llm setup as below:
    # Once this works, then you can try the DocChatAgent
    #
    # agent = lr.ChatAgent(
    #     lr.ChatAgentConfig(
    #         llm=llm
    #     )
    # )
    #
    # agent.llm_response("What is 3 + 4?")
    #
    # task = lr.Task(agent)
    # verify you can interact with this in a chat loop on cmd line:
    # task.run("Concisely answer some questions")

    config = DocChatAgentConfig(
        # default vecdb is qdrantdb
        # using SentenceTransformers/BAAI/bge-large-en-v1.5 embedding model
        llm=llm_config,
        doc_paths=[
            # can be URLS, file-paths, or Folders.
            # File-types: most web-pages, and local pdf, txt, docx
            "https://arxiv.org/pdf/2312.17238.pdf",
        ],
        system_message="""
        Concisely answer my questions about docs. Start by asking me what I want to know.
        """,
    )

    agent = DocChatAgent(config)
    task = lr.Task(agent)
    task.run()


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/extract/capitals.py">
"""
Extract structured information from a passage using a tool/function.


python3 examples/extract/capitals.py

"""

from typing import List

from rich import print

import langroid as lr
from pydantic import BaseModel


class City(BaseModel):
    name: str
    country: str
    population: int


class CitiesData(BaseModel):
    cities: List[City]


PASSAGE = """
Berlin is the capital of Germany. It has a population of 3,850,809. 
Paris, France's capital, has 2.161 million residents. 
Lisbon is the capital and the largest city of Portugal with the population of 504,718.
"""


class CitiesMessage(lr.agent.ToolMessage):
    """Tool/function to use to extract/present structured capitals info"""

    request: str = "capital_info"
    purpose: str = "Collect information about city <capitals> from a passage"
    capitals: List[CitiesData]

    def handle(self) -> str:
        """Tool handler: Print the info about the capitals.
        Any format errors are intercepted by Langroid and passed to the LLM to fix."""
        print(f"Correctly extracted Capitals Info: {self.capitals}")
        return "DONE"  # terminates task


agent = lr.ChatAgent(
    lr.ChatAgentConfig(
        name="CitiesExtractor",
        use_functions_api=True,
        use_tools=False,
        system_message=f"""
        From the passage below, extract info about city capitals, and present it 
        using the `capital_info` tool/function.
        PASSAGE: {PASSAGE}
        """,
    )
)
# connect the Tool to the Agent, so it can use it to present extracted info
agent.enable_message(CitiesMessage)

# wrap the agent in a task and run it
task = lr.Task(
    agent,
    interactive=False,
)

task.run()
</file>

<file path="examples/extract/extract.py">
"""
Extract structured data from text using function_calling/tools.
Inspired by this W&B example notebook, but goes beyond, i.e. gets slightly
more structured output to include model quality:
https://wandb.ai/darek/llmapps/reports/Using-LLMs-to-Extract-Structured-Data-OpenAI-Function-Calling-in-Action--Vmlldzo0Nzc0MzQ3

Example usage, to use Langroid tool:
python3 examples/basic/extract.py -nc

Use -f option to use OpenAI function calling API instead of Langroid tool.

"""

import json
import textwrap
from typing import List

import typer
from kaggle_text import kaggle_description
from rich import print

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from pydantic import BaseModel
from langroid.utils.configuration import Settings, set_global
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


class MethodQuality(BaseModel):
    name: str
    quality: str


class MethodsList(ToolMessage):
    request: str = "methods_list"
    purpose: str = """
        Make a list of Machine Learning methods and their quality
        """
    methods: List[MethodQuality]
    result: str = ""

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(
                methods=[
                    MethodQuality(name="XGBoost", quality="good"),
                    MethodQuality(name="Random Forest", quality="bad"),
                ],
                result="",
            ),
        ]


class ExtractorAgent(ChatAgent):
    def __init__(self, config: ChatAgentConfig):
        super().__init__(config)

    def methods_list(self, message: MethodsList) -> str:
        print(
            f"""
        DONE! Successfully extracted ML Methods list:
        {message.methods}
        """
        )
        return "\n".join(json.dumps(m.model_dump()) for m in message.methods)


class ExtractorConfig(ChatAgentConfig):
    name: str = "Extractor"
    debug: bool = False
    conversation_mode: bool = True
    cache: bool = True  # cache results
    gpt4: bool = False  # use GPT-4?
    stream: bool = True  # allow streaming where needed
    max_tokens: int = 10000
    use_tools: bool = False
    use_functions_api: bool = True
    llm: OpenAIGPTConfig = OpenAIGPTConfig(
        type="openai",
        chat_model=OpenAIChatModel.GPT4o,
    )


def chat(config: ExtractorConfig) -> None:
    print(
        textwrap.dedent(
            """
        [blue]Welcome to the basic chatbot!
        Enter x or q to quit
        """
        ).strip()
    )
    agent = ExtractorAgent(config)
    agent.enable_message(
        MethodsList,
        use=True,
        handle=True,
        force=True,
    )

    task = Task(
        agent,
        system_message="""
        You are a machine learning engineer analyzing Kaggle competition solutions.
        Your goal is to create a list of Machine Learning methods and their 
        quality, based on the user's description. 
        The "quality" can be "good" or "bad", based on your understanding of the 
        description.
        The methods must be very short names, not long phrases.
        Don't add any methods not mentioned in the solution description.
        Call the methods_list function or Tool to accomplish this.
        """,
    )
    task.run(kaggle_description)


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    fn_api: bool = typer.Option(False, "--fn_api", "-f", help="use functions api"),
) -> None:
    config = ExtractorConfig(
        use_functions_api=fn_api,
        use_tools=not fn_api,
    )
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
            cache_type="redis",
        )
    )
    chat(config)


if __name__ == "__main__":
    app()
</file>

<file path="examples/extract/job_listing.txt">
Advantest develops Semiconductor and Component Test Systems. Your job includes improving the C++ development experience for our software engineers developing the software controlling the V93000 test system. No domain knowledge in semiconductor testing is required.

- We are moving the build system from Ant, SCons and Make to Bazel.

- We will be writing AST transformations (Clang AST) to automatically replace dated types and operations with their C++ standard library counterparts.

- We plan to move all code to C++20 modules.

Your impact:

- Contribute to our efforts to move the build of our C++ and Java code from legacy build systems to Bazel.

- Contribute to our efforts to use `clangd` as the LSP-provider of our IDE, Eclipse.

- Implement and run automated transformations to free our code from legacy cruft.

- Provide first-level support to our developers facing issues with the build system.

The skills you will demonstrate:

- Required:

  - Strong knowledge of the Linux environment (Shell, Python, Ruby, Git, RPM-based packaging, GDB, LLDB, util-linux, coreutils, iproute2, bcc, perf, ...)

  - Knowledge of the C++ build process (compiling, linking, build systems).
- Preferred:
  - Knowledge of the LLVM/Clang ecosystem, particularly libAST and clang-tidy internals

  - Development or maintenance of a SCons-, CMake- or Bazel-based build system.
Technologies: C++17 and newer; Linux (RHEL7/RHEL9) only; Bazel; Clang AST (clang-tidy); GNU Make; SCons; Ant; Eclipse CDT; Java
</file>

<file path="examples/extract/kaggle_text.py">
kaggle_description = """
    While it's universally interesting to understand what methods were used by the 
    top participants (especially in this contest where there are some large gaps in 
    AUC at the top), I suspect that many others who participated also have clever 
    methods or insights.  While we wait for the top finishers to post on "No Free 
    Hunch", I thought it would be interesting to hear from anyone else who might wish 
    to share.  Many of the models are quite good and would produce better results 
    than the methods used by persons in industry.      

    My results (#15):

    Overall method: 

    randomForest() in R, 199 trees, min node size of 25, default setting for other 
    values 

    Sampling: 

    Used 10% of the training dataset to train the randomForest.  Also included any 
    data points that were within 500ms of a state change (where isalert shifted from 
    1 to 0 or vice-versa).  About 110,000 rows total.  

    Data Transformations: 

Tossed out correlated variables, such as p7 (inverse correlation with p6) and p4 (
inverse correlation with p3) 
Transformed p3 into an element of ["High", "Mid", "Low"] based on the probability of 
being alert.  Where p3 is an even multiple of 100, the probability of being alert is 
systematically higher.  Where "p3 mod 100" is 84, 16, or 32, there is also a greater 
chance of being alert ("Mid").  Call everything else "Low".   
The histogram of p5 clearly shows a bimodal distribution.  Transformed p5 into a 1/0 
indicator variable with a breakpoint at p5=0.1750. 
Transformed e7 and e8 to lump together all buckets greater than or equal to 4.
Transformed v11 into 20-tiles to convert strangely shaped distribution into a 
discrete variable. 

Tried and Denied:

Lagging values
Moving average


Color Commentary:

RandomForest's ability to "fit" the training data presented was very strong.  
However, the out-of-bucket (OOB) error rate, as reported by R, was highly misleading. 
The OOB error rate could be driven down to the 1-3% range.  However, those models 
produced somewhat worse results on a true out-of-sample validation set.  Keeping 
randomForest tuned to produce OOB error rates of 8-10% produced the best results in 
this case.   

Because many of the training cases are similar, randomForest performed better when 
using just a sample of the overall training data (hence the decision to train on only 
about 110,000 rows).  RandomForest also under-performed when the default nodesize (
either 1 or 5) was used.  The explicit adjustment of nodesize to other values, 
such as 10, 25, and 50, produced noticeably different error rates on true 
out-of-sample data.   
"""
</file>

<file path="examples/extract/lease.html">
<DOCUMENT>
<TYPE>EX-10
<SEQUENCE>2
<FILENAME>elmonteleaseforfiling.htm
<DESCRIPTION>MATERIAL CONTRACT
<TEXT>
<!doctype html public "-//IETF//DTD HTML//EN">
<HTML>
<HEAD>
<TITLE>SAMPLE COMMERCIAL LEASE AGREEMENT</TITLE>
<META NAME="author" CONTENT="Stephen Haas">
<META NAME="date" CONTENT="12/05/2013">
</HEAD>
<BODY style="margin-top:0;font-family:Times New Roman; font-size:10pt; color:#000000">
<DIV style="width:576px"><P style="margin:0px; font-size:12pt" align=center><B>COMMERCIAL LEASE AGREEMENT</B></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">THIS LEASE AGREEMENT is made and entered into on December 1, 2013, by and between Temple CB, LLC, whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as &quot;Landlord&quot;), and Okra Energy, Inc., whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as &quot;Tenant&quot;). </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE I - GRANT OF LEASE</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Landlord, in consideration of the rents to be paid and the covenants and agreements to be performed and observed by the Tenant, does hereby lease to the Tenant and the Tenant does hereby lease and take from the Landlord the property described in Exhibit &quot;A&quot; attached hereto and by reference made a part hereof (the &quot;Leased Premises&quot;), together with, as part of the parcel, all improvements located thereon.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE II - LEASE TERM</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. &nbsp;Term of Lease. &nbsp;The term of this Lease shall begin on the Commencement Date, as defined in Section 2 of this Article II, and shall terminate on May 31, 2020 (&quot;the Termination Date&quot;); provided, however, that at the option of Tenant, Tenant may renew this Lease for five additional successive one- year terms at a Monthly Rent of $100,000 per month, provided that notice of such renewal is given in writing no less than 120 days prior to the Termination Date or the expiration of any one-year renewal term. Tenant may at any time cancel this Lease and terminate all of its obligations hereunder by the payment of $300,000, plus all other amounts then due under this Lease.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Commencement Date. The &quot;Commencement Date&quot; shall mean &nbsp;December 1, 2013. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE III - EXTENSIONS </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">The parties hereto may elect to extend this Agreement upon such terms and conditions as may be agreed upon in writing and signed by the parties at the time of any such extension.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE IV - DETERMINATION OF RENT</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. Monthly Rent: The Tenant agrees to pay the Landlord and the Landlord agrees to accept, during the term hereof, at such place as the Landlord shall from time to time direct by notice to the Tenant, monthly rent set forth in the following table:</P>
<P style="margin:0px"><BR></P>
<P style="margin-top:0px; margin-bottom:-2px; width:432px; font-size:12pt; float:left">Initial Period of December 1, 2013 to May 31, 2014:</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">$ 0</P>
<P style="margin-top:0px; margin-bottom:-2px; width:432px; font-size:12pt; clear:left; float:left">June 1, 2014 to May 31, 2015:</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">$ 30,000</P>
<P style="margin-top:0px; margin-bottom:-2px; width:432px; font-size:12pt; clear:left; float:left">June 1, 2015 to May 31, 2016:</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">$ 40,000</P>
<P style="margin-top:0px; margin-bottom:-2px; width:432px; font-size:12pt; clear:left; float:left">June 1, 2016 to May 31, 2017:</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">$ 50,000</P>
<P style="margin-top:0px; margin-bottom:-2px; width:432px; font-size:12pt; clear:left; float:left">June 1, 2017 to May 31, 2018:</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">$ 60,000</P>
<P style="margin-top:0px; margin-bottom:-2px; width:432px; font-size:12pt; clear:left; float:left">June 1, 2019 to May 31, 2020:</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">$ 70,000</P>
<P style="margin:0px; clear:left"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Late Fee. &nbsp;A late fee in the amount of 5% of the Monthly Rent shall be assessed if payment is not postmarked or received by Landlord on or before the tenth day of each month. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE V - SECURITY DEPOSIT</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">The Tenant has deposited with the Landlord the sum of Twenty Thousand Dollars ($20,000.00) as security for the full and faithful performance by the Tenant of all the terms of this lease required to be performed by the Tenant. Such sum shall be returned to the Tenant after the expiration of this lease, provided the Tenant has fully and faithfully carried out all of its terms. In the event of a bona fide sale of the property of which the leased premises are a part, the Landlord shall have the right to transfer the security to the purchaser to be held under the terms of this lease, and the Landlord shall be released from all liability for the return of such security to the Tenant. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE VI - TAXES</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. &nbsp;Personal Property Taxes. &nbsp;The Tenant shall be liable for all taxes levied against any leasehold interest of the Tenant or personal property and trade fixtures owned or placed by the Tenant in the Leased Premises. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Real Estate Taxes. &nbsp;During the continuance of this lease Landlord shall deliver to Tenant a copy of any real estate taxes and assessments against the Leased Property. From and after the Commencement Date, the Tenant shall pay to Landlord not later than twenty-one (21) days after the day on which the same may become initially due, all real estate taxes and assessments applicable to the Leased Premises, together with any interest and penalties lawfully imposed thereon as a result of Tenant's late payment thereof, which shall be levied upon the Leased Premises during the term of this Lease. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Contest of Taxes. &nbsp;The Tenant, at its own cost and expense, may, if it shall in good faith so desire, contest by appropriate proceedings the amount of any personal or real property tax. The Tenant may, if it shall so desire, endeavor at any time or times, by appropriate proceedings, to obtain a reduction in the assessed valuation of the Leased Premises for tax purposes. In any such event, if the Landlord agrees, at the request of the Tenant, to join with the Tenant at Tenant's expense in said proceedings and the Landlord agrees to sign and deliver such papers and instruments as may be necessary to prosecute such proceedings, the Tenant shall have the right to contest the amount of any such tax and the Tenant shall have the right to withhold payment of any such tax, if the statute under which the Tenant is contesting such tax so permits. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 4. &nbsp;Payment of Ordinary Assessments. &nbsp;The Tenant shall pay all assessments, ordinary and extraordinary, attributable to or against the Leased Premises not later than twenty-one (21) days after the day on which the same became initially due. The Tenant may take the benefit of any law allowing assessments to be paid in installments and in such event the Tenant shall only be liable for such installments of assessments due during the term hereof. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 5. &nbsp;Changes in Method of Taxation. &nbsp;Landlord and Tenant further agree that if at any time during the term of this Lease, the present method of taxation or assessment of real estate shall be changed so that the whole or any part of the real estate taxes, assessment or governmental impositions now levied, assessed or imposed on the Leased Premises shall, in lieu thereof, be assessed, levied, or imposed wholly or in part, as a capital levy or otherwise upon the rents reserved herein or any part thereof, or as a tax, corporation franchise tax, assessment, levy or charge, or any part thereof, measured by or based, in whole or in part, upon the Leased Premises or on the rents derived therefrom and imposed upon the Landlord, then the Tenant shall pay all such taxes, assessments, levies, impositions, or charges. &nbsp;Nothing contained in this Lease shall require the Tenant to pay an estate, inheritance, succession, capital levy, corporate franchise, gross receipts, transfer or income tax of the Landlord, nor shall any of the same be deemed real estate taxes as defined herein unless the same be imposed in lieu of the real estate taxes. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE VII - CONSTRUCTION AND COMPLETION</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Improvements by Tenant. &nbsp;Tenant may have prepared plans and specifications for the construction of improvements, and, if so, such plans and specifications are attached hereto as Exhibit &quot;B&quot; and incorporated herein by reference. Tenant shall obtain all certificates, permits, licenses and other authorizations of governmental bodies or authorities which are necessary to permit the construction of the improvements on the demised premises and shall keep the same in full force and effect at Tenant's cost. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Tenant shall negotiate, let and supervise all contracts for the furnishing of services, labor, and materials for the construction of the improvements on the demised premises at its cost. All such contracts shall require the contracting party to guarantee performance and all workmanship and materials installed by it for a period of one year following the date of completion of construction. &nbsp;Tenant shall cause all contracts to be fully and completely performed in a good and workmanlike manner, all to the effect that the improvements shall be fully and completely constructed and installed in accordance with good engineering and construction practice. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">During the course of construction, Tenant shall, at its cost, keep in full force and effect a policy of builder's risk and liability insurance in a sum equal, from time to time, to three times the amount expended for construction of the improvements. All risk of loss or damage to the improvements during the course of construction shall be on Tenant with the proceeds from insurance thereon payable to Landlord. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Upon completion of construction, Tenant shall, at its cost, obtain an occupancy permit and all other permits or licenses necessary for the occupancy of the improvements and the operation of the same as set out herein and shall keep the same in force. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Nothing herein shall alter the intent of the parties that Tenant shall be fully and completely responsible for all aspects pertaining to the construction of the improvements of the demised premises and for the payment of all costs associated therewith. Landlord shall be under no duty to investigate or verify Tenant's compliance with the provision herein. Moreover, neither Tenant nor any third party may construe the permission granted Tenant hereunder to create any responsibility on the part of the Landlord to pay for any improvements, alterations or repairs occasioned by the Tenant. The Tenant shall keep the property free and clear of all liens and, should the Tenant fail to do so, or to have any liens removed from the property within fourteen (14) days of notification to do so by the Landlord , in addition to all other remedies available to the Landlord , the Tenant shall indemnify and hold the Landlord harmless for all costs and expenses, including attorney's fees, occasioned by the Landlord in having said lien removed from the property; and, such costs and expenses shall be billed to the Tenant monthly and shall be payable by the Tenant with that month's regular monthly rental as additional reimburseable expenses to the Landlord by the Tenant. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Utilities. &nbsp;Tenant shall pay for all water, sanitation, sewer, electricity, light, heat, gas, power, fuel, janitorial, and other services incident to Tenant's use of the Leased Premises, whether or not the cost thereof be a charge or imposition against the Leased Premises. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE VIII - OBLIGATIONS FOR REPAIRS </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Landlord's Repairs. &nbsp;Subject to any provisions herein to the contrary, and except for maintenance or replacement necessitated as the result of the act or omission of sublessees, licensees or contractors, the Landlord shall be required to repair only defects, deficiencies, deviations or failures of materials or workmanship in the building. The Landlord shall keep the Leased Premises free of such defects, deficiencies, deviations or failures during the first twelve (12) months of the term hereof. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Tenant's Repairs. &nbsp;The Tenant shall repair and maintain the Leased Premises in good order and condition, except for reasonable wear and tear, the repairs required of Landlord pursuant hereto, and maintenance or replacement necessitated as the result of the act or omission or negligence of the Landlord, its employees, agents, or contractors. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Requirements of the Law. &nbsp;The Tenant agrees that if any federal, state or municipal government or any department or division thereof shall condemn the Leased Premises or any part thereof as not in conformity with the laws and regulations relating to the construction thereof as of the commencement date with respect to conditions latent or otherwise which existed on the Commencement Date, or, with respect to items which are the Landlord's duty to repair pursuant to Section 1 and 3 of this Article; and such federal, state or municipal government or any other department or division thereof, has ordered or required, or shall hereafter order or require, any alterations or repairs thereof or installations and repairs as may be necessary to comply with such laws, orders or requirements (the validity of which the Tenant shall be entitled to contest); and if by reason of such laws, orders or the work done by the Landlord in connection therewith, the Tenant is deprived of the use of the Leased Premises, the rent shall be abated or adjusted, as the case may be, in proportion to that time during which, and to that portion of the Leased Premises of which, the Tenant shall shall be deprived as a result thereof, and the Landlord shall be obligated to make such repairs, alterations or modifications at Landlord's expense. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">All such rebuilding, altering, installing and repairing shall be done in accordance with Plans and Specifications approved by the Tenant, which approval shall not be unreasonably withheld. If, however, such condemnation, law, order or requirement, as in this Article set forth, shall be with respect to an item which shall be the Tenant's obligation to repair pursuant to Section 2 of this Article VII or with respect to Tenant's own costs and expenses, no abatement or adjustment of rent shall be granted; provided, however, that Tenant shall also be entitled to contest the validity thereof. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 4. &nbsp;Tenant's Alterations. &nbsp;The Tenant shall have the right, at its sole expense, from time to time, to redecorate the Leased Premises and to make such non-structural alterations and changes in such parts thereof as the Tenant shall deem expedient or necessary for its purposes; provided, however, that such alterations and changes shall neither impair the structural soundness nor diminish the value of the Leased Premises. The Tenant may make structural alterations and additions to the Leased Premises provided that Tenant has first obtained the consent thereto of the Landlord in writing. The Landlord agrees that it shall not withhold such consent unreasonably. The Landlord shall execute and deliver upon the request of the Tenant such instrument or instruments embodying the approval of the Landlord which may be required by the public or quasi public authority for the purpose of obtaining any licenses or permits for the making of such alterations, changes and/or installations in, to or upon the Leased Premises and the Tenant agrees to pay for such licenses or permits. &nbsp;The parties understand that a portion of the Leased Premises requires environmental remediation, and the Tenant anticipates that it will undertake such remediation and will be responsible therefore as if it were a structural alteration or addition set forth above.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 5. &nbsp;Permits and Expenses. &nbsp;Each party agrees that it will procure all necessary permits for making any repairs, alterations, or other improvements for installations, when applicable. Each Party hereto shall give written notice to the other party of any repairs required of the other pursuant to the provisions of this Article and the party responsible for said repairs agrees promptly to commence such repairs and to prosecute the same to completion diligently, subject, however, to the delays occasioned by events beyond the control of such party. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Each party agrees to pay promptly when due the entire cost of any work done by it upon the Leased Premises so that the Leased Premises at all times shall be free of liens for labor and materials. &nbsp;Each party further agrees to hold harmless and indemnify the other party from and against any and all injury, loss, claims or damage to any person or property occasioned by or arising out of the doing of any such work by such party or its employees, agents or contractors. Each party further agrees that in doing such work that it will employ materials of good quality and comply with all governmental requirements, and perform such work in a good and workmanlike manner.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE IX - TENANT'S COVENANTS </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. Tenant's Covenants. &nbsp;Tenant covenants and agrees as follows: </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">a. &nbsp;To procure any licenses and permits required for any use made of the Leased Premises by Tenant, and upon the expiration or termination of this Lease, to remove its goods and effects and those of all persons claiming under it, and to yield up peaceably to Landlord the Leased Premises in good order, repair and condition in all respects; excepting only damage by fire and casualty covered by Tenant's insurance coverage, structural repairs (unless Tenant is obligated to make such repairs hereunder) and reasonable wear and tear; </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">b. &nbsp;To permit Landlord and its agents to examine the Leased Premises at reasonable times and to show the Leased Premises to prospective purchasers of the Building and to provide Landlord, if not already available, with a set of keys for the purpose of said examination, provided that Landlord shall not thereby unreasonably interfere with the conduct of Tenant's business; </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">c. &nbsp;To permit Landlord to enter the Leased Premises to inspect such repairs, improvements, alterations or additions thereto as may be required under the provisions of this Lease. If, as a result of such repairs, improvements, alterations, or additions, Tenant is deprived of the use of the Leased Premises, the rent shall be abated or adjusted, as the case may be, in proportion to that time during which, and to that portion of the Leased Premises of which, Tenant shall be deprived as a result thereof. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE X - INDEMNITY BY TENANT</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. Indemnity and Public Liability. &nbsp;The Tenant shall save Landlord harmless and indemnify Landlord from all injury, loss, claims or damage to any person or property while on the Leased Premises, unless caused by the willful acts or omissions or gross negligence of Landlord, its employees, agents, licensees or contractors. Tenant shall maintain, with respect to the Leased Premises, public liability insurance with limits of not less than one million dollars for injury or death from one accident and $250,000.00 property damage insurance, insuring Landlord and Tenant against injury to persons or damage to property on or about the Leased Premises. A copy of the policy or a certificate of insurance shall be delivered to Landlord on or before the commencement date and no such policy shall be cancellable without ten (10) days prior written notice to Landlord. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XI - USE OF PROPERTY BY TENANT</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Use. &nbsp;The Leased Premises may be occupied and used by Tenant exclusively for warehouse and power generation . </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Nothing herein shall give Tenant the right to use the property for any other purpose or to sublease, assign, or license the use of the property to any sublessee, assignee, or licensee, which or who shall use the property for any other use. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XII - SIGNAGE</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. &nbsp;Exterior Signs. &nbsp;Tenant shall have the right, at its sole risk and expense and in conformity with applicable laws and ordinances, to erect and thereafter, to repair or replace, if it shall so elect signs on any portion of the Leased Premises, providing that Tenant shall remove any such signs upon termination of this lease, and repair all damage occasioned thereby to the Leased Premises.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Interior Signs. &nbsp;Tenant shall have the right, at its sole risk and expense and in conformity with applicable laws and ordinances, to erect, maintain, place and install its usual and customary signs and fixtures in the interior of the Leased Premises. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XIII - INSURANCE</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Insurance Proceeds. &nbsp;In the event of any damage to or destruction of the Leased Premises, Tenant shall adjust the loss and settle all claims with the insurance companies issuing such policies. The parties hereto do irrevocably assign the proceeds from such insurance policies for the purposes hereinafter stated to any institutional first mortgagee or to Landlord and Tenant jointly, if no institutional first mortgagee then holds an interest in the Leased Premises. All proceeds of said insurance shall be paid into a trust fund under the control of any institutional first mortgagee, or of Landlord and Tenant if no institutional first mortgagee then holds an interest in the Leased Premises, for repair, restoration, rebuilding or replacement, or any combination thereof, of the Leased Premises or of the improvements in the Leased Premises. In case of such damage or destruction, Landlord shall be entitled to make withdrawals from such trust fund, from time to time, upon presentation of: </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">a. &nbsp;bills for labor and materials expended in repair, restoration, rebuilding or replacement, or any combination thereof; </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">b. &nbsp;Landlord's sworn statement that such labor and materials for which payment is being made have been furnished or delivered on site; and</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">c. &nbsp;the certificate of a supervising architect (selected by Landlord and Tenant and approved by an institutional first mortgagee, if any, whose fees will be paid out of said insurance proceeds) certifying that the work being paid for has been completed in accordance with the Plans and Specifications previously approved by Landlord , Tenant and any institutional first mortgagee in a first class, good and workmanlike manner and in accordance with all pertinent governmental requirements. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Any insurance proceeds in excess of such proceeds as shall be necessary for such repair, restoration, rebuilding, replacement or any combination thereof shall be the sole property of Landlord subject to any rights therein of Landlord's mortgagee, and if the proceeds necessary for such repair, restoration, rebuilding or replacement, or any combination thereof shall be inadequate to pay the cost thereof, Tenant shall suffer the deficiency.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Subrogation. &nbsp;Landlord and Tenant hereby release each other, to the extent of the insurance coverage provided hereunder, from any and all liability or responsibility (to the other or anyone claiming through or under the other by way of subrogation or otherwise) for any loss to or damage of property covered by the fire and extended coverage insurance policies insuring the Leased Premises and any of Tenant's property, even if such loss or damage shall have been caused by the fault or negligence of the other party. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Contribution. Tenant shall reimburse Landlord for all insurance premiums connected with or applicable to the Leased Premises for whatever insurance policy the Landlord , at its sole and exclusive option, should select.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XIV - DAMAGE TO DEMISED PREMISES</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Abatement or Adjustment of Rent. &nbsp;If the whole or any part of the Leased Premises shall be damaged or destroyed by fire or other casualty after the execution of this Lease and before the termination hereof, then in every case the rent reserved in Article IV herein and other charges, if any, shall be abated or adjusted, as the case may be, in proportion to that portion of the Leased Premises of which Tenant shall be deprived on account of such damage or destruction and the work of repair, restoration, rebuilding, or replacement or any combination thereof, of the improvements so damaged or destroyed, shall in no way be construed by any person to effect any reduction of sums or proceeds payable under any rent insurance policy. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Repairs and Restoration. &nbsp;Landlord agrees that in the event of the damage or destruction of the Leased Premises, Landlord forthwith shall proceed to repair, restore, replace or rebuild the Leased Premises (excluding Tenant's leasehold improvements), to substantially the condition in which the same were immediately prior to such damage or destruction. The Landlord thereafter shall diligently prosecute said work to completion without delay or interruption except for events beyond the reasonable control of Landlord . Notwithstanding the foregoing, if Landlord does not either obtain a building permit within ninety (90) days of the date of such damage or destruction, or complete such repairs, rebuilding or restoration and comply with conditions (a), (b) and (c) in Section 1 of Article XIII within nine (9) months of such damage or destruction, then Tenant may at any time thereafter cancel and terminate this Lease by sending ninety (90) days written notice thereof to Landlord , or, in the alternative, Tenant may, during said ninety (90) day period, apply for the same and Landlord shall cooperate with Tenant in Tenant's application. Notwithstanding the foregoing, if such damage or destruction shall occur during the last year of the term of this Lease, or during any renewal term, and shall amount to twenty-five (25%) percent or more of the replacement cost, (exclusive of the land and foundations), this Lease, except as hereinafter provided in Section 3 of Article XV, may be terminated at the election of either Landlord or Tenant, provided that notice of such election shall be sent by the party so electing to the other within thirty (30) days after the occurrence of such damage or destruction. Upon termination, as aforesaid, by either party hereto, this Lease and the term thereof shall cease and come to an end, any unearned rent or other charges paid in advance by Tenant shall be refunded to Tenant, and the parties shall be released hereunder, each to the other, from all liability and obligations hereunder thereafter arising. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XV - CONDEMNATION </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Total Taking. &nbsp;If, after the execution of this Lease and prior to the expiration of the term hereof, the whole of the Leased Premises shall be taken under power of eminent domain by any public or private authority, or conveyed by Landlord to said authority in lieu of such taking, then this Lease and the term hereof shall cease and terminate as of the date when possession of the Leased Premises shall be taken by the taking authority and any unearned rent or other charges, if any, paid in advance, shall be refunded to Tenant.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Partial Taking. &nbsp;If, after the execution of this Lease and prior to the expiration of the term hereof, any public or private authority shall, under the power of eminent domain, take, or Landlord shall convey to said authority in lieu of such taking, property which results in a reduction by fifteen (15%) percent or more of the area in the Leased Premises, or of a portion of the Leased Premises that substantially interrupts or substantially obstructs the conducting of business on the Leased Premises; then Tenant may, at its election, terminate this Lease by giving Landlord notice of the exercise of Tenant's election within thirty (30) days after Tenant shall receive notice of such taking. In the event of termination by Tenant under the provisions of Section 1 of this Article XV, this Lease and the term hereof shall cease and terminate as of the date when possession shall be taken by the appropriate authority of that portion of the Entire Property that results in one of the above takings, and any unearned rent or other charges, if any, paid in advance by Tenant shall be refunded to Tenant.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Restoration. &nbsp;In the event of a taking in respect of which Tenant shall not have the right to elect to terminate this Lease or, having such right, shall not elect to terminate this Lease, this Lease and the term thereof shall continue in full force and effect and Landlord , at Landlord's sole cost and expense, forthwith shall restore the remaining portions of the Leased Premises, including any and all improvements made theretofore to an architectural whole in substantially the same condition that the same were in prior to such taking. A just proportion of the rent reserved herein and any other charges payable by Tenant hereunder, according to the nature and extent of the injury to the Leased Premises and to Tenant's business, shall be suspended or abated until the completion of such restoration and thereafter the rent and any other charges shall be reduced in proportion to the square footage of the Leased Premises remaining after such taking.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 4. &nbsp;The Award. &nbsp;All compensation awarded for any taking, whether for the whole or a portion of the Leased Premises, shall be the sole property of the Landlord whether such compensation shall be awarded for diminution in the value of, or loss of, the leasehold or for diminution in the value of, or loss of, the fee in the Leased Premises, or otherwise. The Tenant hereby assigns to Landlord all of Tenant's right and title to and interest in any and all such compensation. However, the Landlord shall not be entitled to and Tenant shall have the sole right to make its independent claim for and retain any portion of any award made by the appropriating authority directly to Tenant for loss of business, or damage to or depreciation of, and cost of removal of fixtures, personally and improvements installed in the Leased Premises by, or at the expense of Tenant, and to any other award made by the appropriating authority directly to Tenant. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 5. &nbsp;Release. &nbsp;In the event of any termination of this Lease as the result of the provisions of this Article XV, the parties, effective as of such termination, shall be released, each to the other, from all liability and obligations thereafter arising under this lease. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XVI - DEFAULT</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Landlord's Remedies. In the event that: </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">a. &nbsp;Tenant shall on three or more occasions be in default in the payment of rent or other charges herein required to be paid by Tenant (default herein being defined as payment received by Landlord ten or more days subsequent to the due date), regardless of whether or not such default has occurred on consecutive or non-consecutive months; or </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">b. &nbsp;Tenant has caused a lien to be filed against the Landlord's property and said lien is not removed within thirty (30) days of recordation thereof; or </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">c. &nbsp;Tenant shall default in the observance or performance of any of the covenants and agreements required to be performed and observed by Tenant hereunder for a period of thirty (30) days after notice to Tenant in writing of such default (or if such default shall reasonably take more than thirty (30) days to cure, Tenant shall not have commenced the same within the thirty (30) days and diligently prosecuted the same to completion); or </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">d. &nbsp;Sixty (60) days have elapsed after the commencement of any proceeding by or against Tenant, whether by the filing of a petition or otherwise, seeking any reorganization, arrangement, composition, readjustment, liquidation, dissolution or similar relief under the present or future Federal Bankruptcy Act or any other present or future applicable federal, state or other statute or law, whereby such proceeding shall not have been dismissed (provided, however, that the non-dismissal of any such proceeding shall not be a default hereunder so long as all of Tenant's covenants and obligations hereunder are being performed by or on behalf of Tenant); then Landlord shall be entitled to its election (unless Tenant shall cure such default prior to such election), to exercise concurrently or successively, any one or more of the following rights: </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">I. &nbsp;Terminate this Lease by giving Tenant notice of termination, in which event this Lease shall expire and terminate on the date specified in such notice of termination, with the same force and effect as though the date so specified were the date herein originally fixed as the termination date of the term of this Lease, and all rights of Tenant under this Lease and in and to the Premises shall expire and terminate, and Tenant shall remain liable for all obligations under this Lease arising up to the date of such termination, and Tenant shall surrender the Premises to Landlord on the date specified in such notice; or </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ii. &nbsp;Terminate this Lease as provided herein and recover from Tenant all damages Landlord may incur by reason of Tenant's default, including, without limitation, a sum which, at the date of such termination, represents the then value of the excess, if any, of (a) the Minimum Rent, Percentage Rent, Taxes and all other sums which would have been payable hereunder by Tenant for the period commencing with the day following the date of such termination and ending with the date herein before set for the expiration of the full term hereby granted, over (b) the aggregate reasonable rental value of the Premises for the same period, all of which excess sum shall be deemed immediately due and payable; or </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">iii. &nbsp;Without terminating this Lease, declare immediately due and payable all Minimum Rent, Taxes, and other rents and amounts due and coming due under this Lease for the entire remaining term hereof, together with all other amounts previously due, at once; provided, however, that such payment shall not be deemed a penalty or liquidated damages but shall merely constitute payment in advance of rent for the remainder of said term. Upon making such payment, Tenant shall be entitled to receive from Landlord all rents received by Landlord from other assignees, tenants, and subtenants on account of said Premises during the term of this Lease, provided that the monies to which tenant shall so become entitled shall in no event exceed the entire amount actually paid by Tenant to Landlord pursuant to the preceding sentence less all costs, expenses and attorney's fees of Landlord incurred in connection with the reletting of the Premises; or</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">iv. &nbsp;Without terminating this Lease, and with or without notice to Tenant, Landlord may in its own name but as agent for Tenant enter into and upon and take possession of the Premises or any part thereof, and, at landlord's option, remove persons and property therefrom, and such property, if any, may be removed and stored in a warehouse or elsewhere at the cost of, and for the account of Tenant, all without being deemed guilty of trespass or becoming liable for any loss or damage which may be occasioned thereby, and Landlord may rent the Premises or any portion thereof as the agent of Tenant with or without advertisement, and by private negotiations and for any term upon such terms and conditions as Landlord may deem necessary or desirable in order to relet the Premises. Landlord shall in no way be responsible or liable for any rental concessions or any failure to rent the Premises or any part thereof, or for any failure to collect any rent due upon such reletting. Upon such reletting, all rentals received by Landlord from such reletting shall be applied: first, to the payment of any indebtedness (other than any rent due hereunder) from Tenant to Landlord; second, to the payment of any costs and expenses of such reletting,</P>
<P style="margin:0px; font-size:12pt">including, without limitation, brokerage fees and attorney's fees and costs of alterations and repairs; third, to the payment of rent and other charges then due and unpaid hereunder; and the residue, if any shall be held by Landlord to the extent of and for application in payment of future rent as the same may become due and payable hereunder. In reletting the Premises as aforesaid, Landlord may grant rent concessions and Tenant shall not be credited therefor. If such rentals received from such reletting shall at any time or from time to time be less than sufficient to pay to Landlord the entire sums then due from Tenant hereunder, Tenant shall pay any such deficiency to Landlord. Such deficiency shall, at Landlord's option, be calculated and paid monthly. No such reletting shall be construed as an election by Landlord to terminate this Lease unless a written notice of such election has been given to Tenant by Landlord. Notwithstanding any such reletting without termination, Landlord may at any time thereafter elect to terminate this Lease for any such previous default provided same has not been cured; or</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">v. &nbsp;Without liability to Tenant or any other party and without constituting a constructive or actual eviction, suspend or discontinue furnishing or rendering to Tenant any property, material, labor, Utilities or other service, whether Landlord is obligated to furnish or render the same, so long as Tenant is in default under this Lease; or</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">vi. &nbsp;Allow the Premises to remain unoccupied and collect rent from Tenant as it comes due; or </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">vii. &nbsp;Foreclose the security interest described herein, including the immediate taking of possession of all property on or in the Premises; or</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">viii. &nbsp;Pursue such other remedies as are available at law or equity. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">e. &nbsp;Landlord's pursuit of any remedy of remedies, including without limitation, any one or more of the remedies stated herein shall not (1) constitute an election of remedies or preclude pursuit of any other remedy or remedies provided in this Lease or any other remedy or remedies provided by law or in equity, separately or concurrently or in any combination, or (2) sever as the basis for any claim of constructive eviction, or allow Tenant to withhold any payments under this Lease.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Landlord's Self Help. &nbsp;If in the performance or observance of any agreement or condition in this Lease contained on its part to be performed or observed and shall not cure such default within thirty (30) days after notice from Landlord specifying the default (or if such default shall reasonably take more than thirty (30) days to cure, shall diligently prosecuted the same to completion), Landlord may, at its option, without waiving any claim for damages for breach of agreement, at any time thereafter cure such default for the account of Tenant, and any amount paid or contractual liability incurred by Landlord in so doing shall be deemed paid or incurred for the account of Tenant and Tenant agrees to reimburse Landlord therefor and save Landlord harmless therefrom. Provided, however, that Landlord may cure any such default as aforesaid prior to the expiration of said waiting period, without notice to Tenant if any emergency situation exists, or after notice to Tenant, if the curing of such default prior to the expiration of said waiting period is reasonably necessary to protect the Leased Premises or Landlord's interest therein, or to prevent injury or damage to persons or property. If Tenant shall fail to reimburse Landlord upon demand for any amount paid for the account of Tenant hereunder, said amount shall be added to and become due as a part of the next payment of rent due and shall for all purposes be deemed and treated as rent hereunder.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Tenant's Self Help. &nbsp;If Landlord shall default in the performance or observance of any agreement or condition in this Lease contained on its part to be performed or observed, and if Landlord shall not cure such default within thirty (30) days after notice from Tenant specifying the default (or, if such default shall reasonably take more than thirty (30) days to cure, and Landlord shall not have commenced the same within the thirty (30) days and diligently prosecuted the same to completion), Tenant may, at its option, without waiving any claim for damages for breach of agreement, at any time thereafter cure such default for the account of Landlord and any amount paid or any contractual liability incurred by Tenant in so doing shall be deemed paid or incurred for the account of Landlord and Landlord shall reimburse Tenant therefor and save Tenant harmless therefrom. Provided, however, that Tenant may cure any such default as aforesaid prior to the expiration of said waiting period, without notice to Landlord if an emergency situation exists, or after notice to Landlord , if the curing of such default prior to the expiration of said waiting period is reasonably necessary to protect the Leased Premises or Tenant's interest therein or to prevent injury or damage to persons or property. &nbsp;If Landlord shall fail to reimburse Tenant upon demand for any amount paid or liability incurred for the account of Landlord hereunder, said amount or liability may be deducted by Tenant from the next or any succeeding payments of rent due hereunder; provided, however, that should said amount or the liability therefor be disputed by Landlord, Landlord may contest its liability or the amount thereof, through arbitration or through a declaratory judgment action and Landlord shall bear the cost of the filing fees therefor.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XVII - TITLE</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. &nbsp;Subordination. &nbsp;Tenant shall, upon the request of Landlord in writing, subordinate this Lease to the lien of any present or future institutional mortgage upon the Leased Premises irrespective of the time of execution or the time of recording of any such mortgage. Provided, however, that as a condition to such subordination, the holder of any such mortgage shall enter first into a written agreement with Tenant in form suitable for recording to the effect that:</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">a. &nbsp;in the event of foreclosure or other action taken under the mortgage by the holder thereof, this Lease and the rights of Tenant hereunder shall not be disturbed but shall continue in full force and effect so long as Tenant shall not be in default hereunder, and </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">b. &nbsp;such holder shall permit insurance proceeds and condemnation proceeds to be used for any restoration and repair required by the provisions of Articles XIII, XIV or XV, respectively. &nbsp;Tenant agrees that if the mortgagee or any person claiming under the mortgagee shall succeed to the interest of Landlord in this Lease, Tenant will recognize said mortgagee or person as its Landlord under the terms of this Lease, provided that said mortgagee or person for the period during which said mortgagee or person respectively shall be in possession of the Leased Premises and thereafter their respective successors in interest shall assume all of the obligations of Landlord hereunder. The word &quot;mortgage&quot;, as used herein includes mortgages, deeds of trust or other similar instruments, and modifications, and extensions thereof. The term &quot;institutional mortgage&quot; as used in this Article XVII means a mortgage securing a loan from a bank (commercial or savings) or trust company, insurance company or pension trust or any other lender institutional in nature and constituting a lien upon the Leased Premises. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Quiet Enjoyment. &nbsp;Landlord covenants and agrees that upon Tenant paying the rent and observing and performing all of the terms, covenants and conditions on Tenant's part to be observed and performed hereunder, that Tenant may peaceably and quietly have, hold, occupy and enjoy the Leased Premises in accordance with the terms of this Lease without hindrance or molestation from Landlord or any persons lawfully claiming through Landlord . </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Zoning and Good Title. &nbsp;Landlord warrants and represents, upon which warranty and representation Tenant has relied in the execution of this Lease, that Landlord is the owner of the Leased Premises, in fee simple absolute, free and clear of all encumbrances, except for the easements, covenants and restrictions of record as of the date of this Lease. Such exceptions shall not impede or interfere with the quiet use and enjoyment of the Leased Premises by Tenant. Landlord further warrants and covenants that this Lease is and shall be a first lien on the Leased Premises, subject only to any Mortgage to which this Lease is subordinate or may become subordinate pursuant to an agreement executed by Tenant, and to such encumbrances as shall be caused by the acts or omissions of Tenant; that Landlord has full right and lawful authority to execute this Lease for the term, in the manner, and upon the conditions and provisions herein contained; that there is no legal impediment to the use of the Leased Premises as set out herein; that the Leased Premises are not subject to any easements, restrictions, zoning ordinances or similar governmental regulations which prevent their use as set out herein; that the Leased Premises presently are zoned for the use contemplated herein and throughout the term of this lease may continue to be so used therefor by virtue of said zoning, under the doctrine of &quot;non-conforming use&quot;, or valid and binding decision of appropriate authority, except, however, that said representation and warranty by Landlord shall not be applicable in the event that Tenant's act or omission shall invalidate the application of said zoning, the doctrine of &quot;non-conforming use&quot; or the valid and binding decision of the appropriate authority. Landlord shall furnish without expense to Tenant, within thirty (30) days after written request therefor by Tenant, a title report covering the Leased Premises showing the condition of title as of the date of such certificate, provided, however, that Landlord's obligation hereunder shall be limited to the furnishing of only one such title report. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 4. &nbsp;Licenses. &nbsp;It shall be the Tenant's responsibility to obtain any and all necessary licenses and the Landlord shall bear no responsibility therefor; the Tenant shall promptly notify Landlord of the fact that it has obtained the necessary licenses in order to prevent any delay to Landlord in commencing construction of the Leased Premises. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XVIII - EXTENSIONS/WAIVERS/DISPUTES</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. &nbsp;Extension Period. &nbsp;Any extension hereof shall be subject to the provisions of Article III hereof. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Holding Over. &nbsp;In the event that Tenant or anyone claiming under Tenant shall continue occupancy of the Leased Premises after the expiration of the term of this Lease or any renewal or extension thereof without any agreement in writing between Landlord and Tenant with respect thereto, such occupancy shall not be deemed to extend or renew the term of the Lease, but such occupancy shall continue as a tenancy at will, from month to month, upon the covenants, provisions and conditions herein contained. The rental shall be the rental in effect during the term of this Lease as extended or renewed, prorated and payable for the period of such occupancy. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Waivers. &nbsp;Failure of either party to complain of any act or omission on the part of the other party, no matter how long the same may continue, shall not be deemed to be a waiver by said party of any of its rights hereunder. No waiver by either party at any time, express or implied, of any breach of any provision of this Lease shall be deemed a waiver of a breach of any other provision of this Lease or a consent to any subsequent breach of the same or any other provision. If any action by either party shall require the consent or approval of the other party, the other party's consent to or approval of such action on any one occasion shall not be deemed a consent to or approval of said action on any subsequent occasion or a consent to or approval of any other action on the same or any subsequent occasion. Any and all rights and remedies which either party may have under this Lease or by operation of law, either at law or in equity, upon any breach, shall be distinct, separate and cumulative and shall not be deemed inconsistent with each other, and no one of them, whether exercised by said party or not, shall be deemed to be an exclusion of any other; and any two or more or all of such rights and remedies may be exercised at the same time. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 4. &nbsp;Disputes. &nbsp;It is agreed that, if at any time a dispute shall arise as to any amount or sum of money to be paid by one party to the other under the provisions hereof, the party against whom the obligation to pay the money is asserted shall have the right to make payment &quot;under protest&quot; and such payment shall not be regarded as a voluntary payment and there shall survive the right on the part of the said party to institute suit for the recovery of such sum. If it shall be adjudged that there was no legal obligation on the part of said party to pay such sum or any part thereof, said party shall be entitled to recover such sum or so much thereof as it was not legally required to pay under the provisions of this Lease. If at any time a dispute shall arise between the parties hereto as to any work to be performed by either of them under the provisions hereof, the party against whom the obligation to perform the work is asserted may perform such work and pay the costs thereof &quot;under protest&quot; and the performance of such work shall in no event be regarded as a voluntary performance and shall survive the right on the part of the said party to institute suit for the recovery of the costs of such work. If it shall be adjudged that there was no legal obligation on the part of the said party to perform the same or any part thereof, said party shall be entitled to recover the costs of such work or the cost of so much thereof as said party was not legally required to perform under the provisions of this Lease and the amount so paid by Tenant may be withheld or deducted by Tenant from any rents herein reserved. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 5. &nbsp;Tenant's Right to cure Landlord's Default. &nbsp;In the event that Landlord shall fail, refuse or neglect to pay any mortgages, liens or encumbrances, the judicial sale of which might affect the interest of Tenant hereunder, or shall fail, refuse or neglect to pay any interest due or payable on any such mortgage, lien or encumbrance, Tenant may pay said mortgages, liens or encumbrances, or interest or perform said conditions and charge to Landlord the amount so paid and withhold and deduct from any rents herein reserved such amounts so paid, and any excess over and above the amounts of said rents shall be paid by Landlord to Tenant. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 6. &nbsp;Notices. &nbsp;All notices and other communications authorized or required hereunder shall be in writing and shall be given by mailing the same by certified mail, return receipt requested, postage prepaid, and any such notice or other communication shall be deemed to have been given when received by the party to whom such notice or other communication shall be addressed. If intended for Landlord the same will be mailed to the address herein above set forth or such other address as Landlord may hereafter designate by notice to Tenant, and if intended for Tenant, the same shall be mailed to Tenant at the address herein above set forth, or such other address or addresses as Tenant may hereafter designate by notice to Landlord. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XIX - PROPERTY DAMAGE </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section l. &nbsp;Loss and Damage. &nbsp;Notwithstanding any contrary provisions of this Lease, Landlord shall not be responsible for any loss of or damage to property of Tenant or of others located on the Leased Premises, except where caused by the willful act or omission or negligence of Landlord , or Landlord's agents, employees or contractors, provided, however, that if Tenant shall notify Landlord in writing of repairs which are the responsibility of Landlord under Article VII hereof, and Landlord shall fail to commence and diligently prosecute to completion said repairs promptly after such notice, and if after the giving of such notice and the occurrence of such failure, loss of or damage to Tenant's property shall result from the condition as to which Landlord has been notified, Landlord shall indemnify and hold harmless Tenant from any loss, cost or expense arising therefrom. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Force Majeure. &nbsp;In the event that Landlord or Tenant shall be delayed or hindered in or prevented from the performance of any act other than Tenant's obligation to make payments of rent, additional rent, and other charges required hereunder, by reason of strikes, lockouts, unavailability of materials, failure of power, restrictive governmental laws or regulations, riots, insurrections, the act, failure to act, or default of the other party, war or other reason beyond its control, then performance of such act shall be excused for the period of the delay and the period for the performance of such act shall be extended for a period equivalent to the period of such delay. &nbsp;Notwithstanding the foregoing, lack of funds shall not be deemed to be a cause beyond control of either party. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XX - OPTION TO PURCHASE</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">During the Term of this Lease, Tenant shall have the right to purchase the Leased Premises at any time for a purchase price equal to Three Million Dollars ($3,000,000).</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">ARTICLE XXI - MISCELLANEOUS </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 1. &nbsp;Assignment and Subletting. &nbsp;Under the terms and conditions hereunder, Tenant shall have the absolute right to transfer and assign this lease or to sublet all or any portion of the Leased Premises or to cease operating Tenant's business on the Leased Premises provided that at the time of such assignment or sublease Tenant shall not be in default in the performance and observance of the obligations imposed upon Tenant hereunder. The use of the Leased Premises by such assignee or sublessee shall be expressly limited by and to the provisions of this lease. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 2. &nbsp;Fixtures. &nbsp;All personal property, furnishings and equipment presently and all other trade fixtures installed in or hereafter by or at the expense of Tenant and all additions and/or improvements, exclusive of structural, mechanical, electrical, and plumbing, affixed to the Leased Premises and used in the operation of the Tenant's business made to, in or on the Leased Premises by and at the expense of Tenant and susceptible of being removed from the Leased Premises without damage, unless such damage be repaired by Tenant, shall remain the property of Tenant and Tenant may, but shall not be obligated to, remove the same or any part thereof at any time or times during the term hereof, provided that Tenant, at its sole cost and expense, shall make any repairs occasioned by such removal. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 3. &nbsp;Estoppel Certificates. &nbsp;At any time and from time to time, Landlord and Tenant each agree, upon request in writing from the other, to execute, acknowledge and deliver to the other or to any person designated by the other a statement in writing certifying that the Lease is unmodified and is in full force and effect, or if there have been modifications, that the same is in full force and effect as modified (stating the modifications), that the other party is not in default in the performance of its covenants hereunder, or if there have been such defaults, specifying the same, and the dates to which the rent and other charges have been paid.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 4. &nbsp;Invalidity of Particular Provision. &nbsp;If any term or provision of this Lease or the application hereof to any person or circumstance shall, to any extent, be held invalid or unenforceable, the remainder of this Lease, or the application of such term or provision to persons or circumstances other than those as to which it is held invalid or unenforceable, shall not be affected thereby, and each term and provision of this Lease shall be valid and be enforced to the fullest extent permitted by law. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 5. &nbsp;Captions and Definitions of Parties. &nbsp;The captions of the Sections of this Lease are for convenience only and are not a part of this Lease and do not in any way limit or amplify the terms and provisions of this Lease. The word &quot;Landlord&quot; and the pronouns referring thereto, shall mean, where the context so admits or requires, the persons, firm or corporation named herein as Landlord or the mortgagee in possession at any time, of the land and building comprising the Leased Premises. If there is more than one Landlord, the covenants of Landlord shall be the joint and several obligations of each of them, and if Landlord is a partnership, the covenants of Landlord shall be the joint and several obligations of each of the partners and the obligations of the firm. Any pronoun shall be read in the singular or plural and in such gender as the context may require. Except as in this Lease otherwise provided, the terms and provisions of this Lease shall be binding upon and inure to the benefit of the parties hereto and their respective successors and assigns. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Nothing contained herein shall be deemed or construed by the parties hereto nor by any third party as creating the relationship of principal and agent or of partnership or of a joint venture between the parties hereto, it being understood and agreed that neither any provision contained herein, nor any acts of the parties hereto, shall be deemed to create any relationship between the parties hereto other than the relationship of Landlord and Tenant. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 6. &nbsp;Brokerage. &nbsp;No party has acted as, by or through a broker in the effectuation of this Agreement, except as set out hereinafter. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 7. &nbsp;Entire Agreement. &nbsp;This instrument contains the entire and only agreement between the parties, and no oral statements or representations or prior written matter not contained in this instrument shall have any force and effect. This Lease shall not be modified in any way except by a writing executed by both parties. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 8. &nbsp;Governing Law. All matters pertaining to this agreement (including its interpretation, application, validity, performance and breach) in whatever jurisdiction action may be brought, shall be governed by, construed and enforced in accordance with the laws of the State of California. The parties herein waive trial by jury and agree to submit to the personal jurisdiction and venue of a court of subject matter jurisdiction located in Los Angeles County, State of California. &nbsp;In the event that litigation results from or arises out of this Agreement or the performance thereof, the parties agree to reimburse the prevailing party's reasonable attorney's fees, court costs, and all other expenses, whether or not taxable by the court as costs, in addition to any other relief to which the prevailing party may be entitled. In such event, no action shall be entertained by said court or any court of competent jurisdiction if filed more than one year subsequent to the date the cause(s) of action actually accrued regardless of whether damages were otherwise as of said time calculable. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 9. &nbsp;Contractual Procedures. &nbsp;Unless specifically disallowed by law, should litigation arise hereunder, service of process therefor may be obtained through certified mail, return receipt requested; the parties hereto waiving any and all rights they may have to object to the method by which service was perfected. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 10. &nbsp;Extraordinary remedies. &nbsp;To the extent cognizable at law, the parties hereto, in the event of breach and in addition to any and all other remedies available thereto, may obtain injunctive relief, regardless of whether the injured party can demonstrate that no adequate remedy exists at law. </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Section 11. &nbsp;Reliance on Financial Statement. &nbsp;Tenant shall furnish concurrently with the execution of this lease, a financial statement of Tenant prepared by an accountant. Tenant, both in corporate capacity, if applicable, and individually, hereby represents and warrants that all the information contained therein is complete, true, and correct. Tenant understands that Landlord is relying upon the accuracy of the information contained therein. Should there be found to exist any inaccuracy within the financial statement which adversely affects Tenant's financial standing, or should Tenant's financial circumstances materially change, Landlord may demand, as additional security, an amount equal to an additional two (2) months' rent, which additional security shall be subject to all terms and conditions herein, require a fully executed guaranty by a third party acceptable to Landlord, elect to terminate this Lease, or hold Tenant personally and individually liable hereunder.</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">IN WITNESS WHEREOF, the parties hereto have executed this Lease the day and year first above written or have caused this Lease to be executed by their respective officers thereunto duly authorized. </P>
<P style="margin:0px"><BR></P>
<P style="margin-top:0px; margin-bottom:-2px; width:288px; font-size:12pt; float:left">TEMPLE CB, LLC</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">OKRA ENERGY, INC.</P>
<P style="margin:0px; clear:left"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin-top:0px; margin-bottom:-2px; width:288px; font-size:12pt; float:left">/s/ Jay Hooper</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">Jay Hooper</P>
<P style="margin-top:0px; margin-bottom:-2px; width:288px; font-size:12pt; clear:left; float:left">Jay Hooper, Manager &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</P>
<P style="margin:0px; text-indent:-2px; font-size:12pt">Jay Hooper, President</P>
<P style="margin:0px; clear:left"><BR></P>
<P style="line-height:12pt; margin:0px; text-indent:336px; font-family:Arial; font-size:12pt">&nbsp;</P>
<P style="line-height:12pt; margin-top:0px; margin-bottom:-2px; width:240px; font-size:12pt; float:left">STATE OF CALIFORNIA</P>
<P style="line-height:12pt; margin:0px; text-indent:-2px; font-size:12pt" align=justify>}</P>
<P style="line-height:12pt; margin-top:0px; margin-bottom:-2px; text-indent:240px; width:288px; font-size:12pt; clear:left; float:left">}</P>
<P style="line-height:12pt; margin:0px; text-indent:-2px; font-size:12pt" align=justify>ss.</P>
<P style="line-height:12pt; margin:0px; font-size:12pt; clear:left" align=justify>COUNTY OF LOS ANGELES &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;}</P>
<P style="line-height:12pt; margin:0px" align=justify><BR></P>
<P style="line-height:12pt; margin:0px; text-indent:49.467px; font-size:12pt">On this __ day of &nbsp;December, &nbsp;2013, before me, the undersigned, a Notary Public in and for said County and State, residing therein, duly commissioned and sworn, personally appeared &nbsp;Jay Hooper, personally known to me (or proved to me on the basis of satisfactory evidence) to be the person whose name is subscribed to the within instrument and acknowledged to me that he executed the same in his authorized capacity, and that by his signature on the instrument the person, or the entity upon behalf of which the person acted, executed the instrument.</P>
<P style="line-height:12pt; margin:0px"><BR></P>
<P style="line-height:12pt; margin:0px; text-indent:49.467px; font-size:12pt" align=justify>WITNESS my hand and official seal.</P>
<P style="line-height:12pt; margin:0px" align=justify><BR></P>
<P style="line-height:12pt; margin:0px; text-indent:336px; font-size:12pt" align=justify>Signature:</P>
<P style="line-height:12pt; margin:0px" align=justify><BR></P>
<P style="line-height:12pt; margin:0px; text-indent:336px; font-size:12pt" align=justify>Name (typed or printed)</P>
<P style="line-height:12pt; margin:0px" align=justify><BR></P>
<P style="line-height:12pt; margin:0px; text-indent:336px; font-size:12pt" align=justify>My Commission expires:</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt" align=center>EXHIBIT &quot;A&quot; LEGAL DESCRIPTION</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">The following described real property, together with all improvements thereon: </P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Initials: </P>
<P style="margin:0px; font-size:12pt">LANDLORD &nbsp;______________</P>
<P style="margin:0px; font-size:12pt">TENANT &nbsp;&nbsp;______________</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">EXHIBIT &quot;B&quot; TENANT PLANS AND SPECIFICATIONS</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
<P style="margin:0px; font-size:12pt">Initials: </P>
<P style="margin:0px; font-size:12pt">LANDLORD &nbsp;_____________</P>
<P style="margin:0px; font-size:12pt">TENANT &nbsp;&nbsp;_____________</P>
<P style="margin:0px; font-size:12pt">&nbsp;</P>
<P style="margin:0px"><BR>
<BR></P>
<P style="margin:0px; font-size:12pt" align=center>1</P>
<P style="margin:0px"><BR></P>
<P style="margin:0px"><BR></P>
</DIV></BODY>
<!-- EDGAR Validation Code: C237D44C -->
</HTML>
</TEXT>
</DOCUMENT>
</file>

<file path="examples/extract/lease.txt">
EX-10 2 elmonteleaseforfiling.htm MATERIAL CONTRACT
COMMERCIAL LEASE AGREEMENT



THIS LEASE AGREEMENT is made and entered into on December 1, 2013, by and between Temple CB, LLC, whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as "Landlord"), and Okra Energy, Inc., whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as "Tenant").



ARTICLE I - GRANT OF LEASE



Landlord, in consideration of the rents to be paid and the covenants and agreements to be performed and observed by the Tenant, does hereby lease to the Tenant and the Tenant does hereby lease and take from the Landlord the property described in Exhibit "A" attached hereto and by reference made a part hereof (the "Leased Premises"), together with, as part of the parcel, all improvements located thereon.



ARTICLE II - LEASE TERM



Section l.  Term of Lease.  The term of this Lease shall begin on the Commencement Date, as defined in Section 2 of this Article II, and shall terminate on May 31, 2020 ("the Termination Date"); provided, however, that at the option of Tenant, Tenant may renew this Lease for five additional successive one- year terms at a Monthly Rent of $100,000 per month, provided that notice of such renewal is given in writing no less than 120 days prior to the Termination Date or the expiration of any one-year renewal term. Tenant may at any time cancel this Lease and terminate all of its obligations hereunder by the payment of $300,000, plus all other amounts then due under this Lease.



Section 2.  Commencement Date. The "Commencement Date" shall mean  December 1, 2013.



ARTICLE III - EXTENSIONS



The parties hereto may elect to extend this Agreement upon such terms and conditions as may be agreed upon in writing and signed by the parties at the time of any such extension.



ARTICLE IV - DETERMINATION OF RENT



Section 1. Monthly Rent: The Tenant agrees to pay the Landlord and the Landlord agrees to accept, during the term hereof, at such place as the Landlord shall from time to time direct by notice to the Tenant, monthly rent set forth in the following table:



Initial Period of December 1, 2013 to May 31, 2014:
$ 0

June 1, 2014 to May 31, 2015:
$ 30,000

June 1, 2015 to May 31, 2016:
$ 40,000

June 1, 2016 to May 31, 2017:
$ 50,000

June 1, 2017 to May 31, 2018:
$ 60,000

June 1, 2019 to May 31, 2020:
$ 70,000



Section 2.  Late Fee.  A late fee in the amount of 5% of the Monthly Rent shall be assessed if payment is not postmarked or received by Landlord on or before the tenth day of each month.



ARTICLE V - SECURITY DEPOSIT



The Tenant has deposited with the Landlord the sum of Twenty Thousand Dollars ($20,000.00) as security for the full and faithful performance by the Tenant of all the terms of this lease required to be performed by the Tenant. Such sum shall be returned to the Tenant after the expiration of this lease, provided the Tenant has fully and faithfully carried out all of its terms. In the event of a bona fide sale of the property of which the leased premises are a part, the Landlord shall have the right to transfer the security to the purchaser to be held under the terms of this lease, and the Landlord shall be released from all liability for the return of such security to the Tenant.



ARTICLE VI - TAXES



Section l.  Personal Property Taxes.  The Tenant shall be liable for all taxes levied against any leasehold interest of the Tenant or personal property and trade fixtures owned or placed by the Tenant in the Leased Premises.



Section 2.  Real Estate Taxes.  During the continuance of this lease Landlord shall deliver to Tenant a copy of any real estate taxes and assessments against the Leased Property. From and after the Commencement Date, the Tenant shall pay to Landlord not later than twenty-one (21) days after the day on which the same may become initially due, all real estate taxes and assessments applicable to the Leased Premises, together with any interest and penalties lawfully imposed thereon as a result of Tenant's late payment thereof, which shall be levied upon the Leased Premises during the term of this Lease.



Section 3.  Contest of Taxes.  The Tenant, at its own cost and expense, may, if it shall in good faith so desire, contest by appropriate proceedings the amount of any personal or real property tax. The Tenant may, if it shall so desire, endeavor at any time or times, by appropriate proceedings, to obtain a reduction in the assessed valuation of the Leased Premises for tax purposes. In any such event, if the Landlord agrees, at the request of the Tenant, to join with the Tenant at Tenant's expense in said proceedings and the Landlord agrees to sign and deliver such papers and instruments as may be necessary to prosecute such proceedings, the Tenant shall have the right to contest the amount of any such tax and the Tenant shall have the right to withhold payment of any such tax, if the statute under which the Tenant is contesting such tax so permits.



Section 4.  Payment of Ordinary Assessments.  The Tenant shall pay all assessments, ordinary and extraordinary, attributable to or against the Leased Premises not later than twenty-one (21) days after the day on which the same became initially due. The Tenant may take the benefit of any law allowing assessments to be paid in installments and in such event the Tenant shall only be liable for such installments of assessments due during the term hereof.
</file>

<file path="examples/extract/least-truncated.txt">
THIS LEASE AGREEMENT is made and entered into on December 1, 2013, by and between Temple CB, LLC, whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as "Landlord"), and Okra Energy, Inc., whose address is 4350 Temple City Boulevard, El Monte, California 91731 (hereinafter referred to as "Tenant").



ARTICLE I - GRANT OF LEASE



Landlord, in consideration of the rents to be paid and the covenants and agreements to be performed and observed by the Tenant, does hereby lease to the Tenant and the Tenant does hereby lease and take from the Landlord the property described in Exhibit "A" attached hereto and by reference made a part hereof (the "Leased Premises"), together with, as part of the parcel, all improvements located thereon.



ARTICLE II - LEASE TERM



Section l.  Term of Lease.  The term of this Lease shall begin on the Commencement Date, as defined in Section 2 of this Article II, and shall terminate on May 31, 2020 ("the Termination Date"); provided, however, that at the option of Tenant, Tenant may renew this Lease for five additional successive one- year terms at a Monthly Rent of $100,000 per month, provided that notice of such renewal is given in writing no less than 120 days prior to the Termination Date or the expiration of any one-year renewal term. Tenant may at any time cancel this Lease and terminate all of its obligations hereunder by the payment of $300,000, plus all other amounts then due under this Lease.



Section 2.  Commencement Date. The "Commencement Date" shall mean  December 1, 2013.



ARTICLE III - EXTENSIONS



The parties hereto may elect to extend this Agreement upon such terms and conditions as may be agreed upon in writing and signed by the parties at the time of any such extension.



ARTICLE IV - DETERMINATION OF RENT



Section 1. Monthly Rent: The Tenant agrees to pay the Landlord and the Landlord agrees to accept, during the term hereof, at such place as the Landlord shall from time to time direct by notice to the Tenant, monthly rent set forth in the following table:



Initial Period of December 1, 2013 to May 31, 2014:
$ 0

June 1, 2014 to May 31, 2015:
$ 30,000

June 1, 2015 to May 31, 2016:
$ 40,000

June 1, 2016 to May 31, 2017:
$ 50,000

June 1, 2017 to May 31, 2018:
$ 60,000

June 1, 2019 to May 31, 2020:
$ 70,000



Section 2.  Late Fee.  A late fee in the amount of 5% of the Monthly Rent shall be assessed if payment is not postmarked or received by Landlord on or before the tenth day of each month.



ARTICLE V - SECURITY DEPOSIT



The Tenant has deposited with the Landlord the sum of Twenty Thousand Dollars ($20,000.00) as security for the full and faithful performance by the Tenant of all the terms of this lease required to be performed by the Tenant. Such sum shall be returned to the Tenant after the expiration of this lease, provided the Tenant has fully and faithfully carried out all of its terms. In the event of a bona fide sale of the property of which the leased premises are a part, the Landlord shall have the right to transfer the security to the purchaser to be held under the terms of this lease, and the Landlord shall be released from all liability for the return of such security to the Tenant.
</file>

<file path="examples/extract/pdf-json-flex.py">
"""
Extract an arbitrary json structure from a pdf via markdown.

1. use Langroid's PDF Parser with `marker` library to
   extract content from (pdf) report in markdown format
2. use Langroid Agent equipped with a structured output tool to extract structured data

Run like this: (drop the -m arg to default to GPT4o)

uv run examples/pdf-json-flex.py -f examples/extract/um-financial-report.pdf \
    -m gemini/gemini-2.0-pro-exp-02-05

NOTES:
- this script uses the `marker` library for parsing PDF content,
and to get that to work with langroid, install langroid with the `marker-pdf` extra,
e.g.
uv pip install "langroid[marker-pdf]"
pip install "langroid[marker-pdf]"

- The structured extracted is very simple, consisting of 3 fields: item, year, and value.
  You may need to adapt it to your needs.
"""

import logging
import os
from typing import List

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ResultTool
from langroid.parsing.document_parser import DocumentParser
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig
from pydantic import BaseModel, ConfigDict

logger = logging.getLogger(__name__)


os.environ["TOKENIZERS_PARALLELISM"] = "false"

"""
Agent to extract structured data from a markdown formatted table.
Typically, this markdown formatted content would come from using a parser
that extracted markdown from a pdf report, e.g. using a Langroid PDF Parser.
"""


class JsonData(BaseModel):
    """Data model for arbitrary nested JSON-like structures.

    This model allows for storing any valid JSON data format, including nested objects,
    arrays, primitives, etc.

    """

    model_config = ConfigDict(extra="allow")  # Allow any extra fields


class FinalResult(ResultTool):
    data: List[JsonData]


class JsonExtractTool(lr.ToolMessage):
    request: str = "json_extract_tool"
    purpose: str = "To present the <json_data> extracted from a document."

    json_data: List[JsonData]

    def handle(self) -> FinalResult:
        return FinalResult(data=self.json_data)


class JsonExtractorConfig(lr.ChatAgentConfig):
    # placeholder
    name: str = "JsonExtractor"


class JsonExtractor(lr.ChatAgent):
    def __init__(self, config: JsonExtractorConfig):
        super().__init__(config)
        self.config: JsonExtractorConfig = config
        self.enable_message(JsonExtractTool)


def display_json_data(data: List[JsonData]) -> None:
    """Display structured JSON data using Rich's JSON pretty printer.

    Args:
        data: List of JsonData objects to display
    """
    from rich.console import Console
    from rich.json import JSON
    from rich.panel import Panel

    console = Console()

    if not data:
        console.print("[bold red]No data found[/bold red]")
        return

    for i, item in enumerate(data):
        # Convert JsonData to dict, filtering out internal attributes
        item_dict = {k: v for k, v in item.__dict__.items() if not k.startswith("__")}
        # Create a panel for each data item with pretty-printed JSON inside
        json_str = JSON.from_data(item_dict)
        console.print(Panel(json_str, title=f"Item {i+1}", border_style="cyan"))

        # Add some spacing between items
        if i < len(data) - 1:
            console.print("")


def make_json_extractor_task(
    llm_config: lm.OpenAIGPTConfig = lm.OpenAIGPTConfig(
        chat_model=lm.OpenAIChatModel.GPT4o,
    )
):
    agent = JsonExtractor(
        JsonExtractorConfig(
            llm=llm_config,
            handle_llm_no_tool=f"You FORGOT to use the TOOL `{JsonExtractTool.name()}`",
            system_message=f"""
            You are an expert at creating (possibly nested) JSON structures
            from markdown documents.
            
            When you receive a markdown-formatted document, your job is to
            extract the data from the document and present it in a structured
            form using the TOOL `{JsonExtractTool.name()}`.
            """,
        )
    )
    # create task specialized to return FinalResult value
    task = lr.Task(agent, interactive=False, single_round=False)[FinalResult]
    return task


def main(
    filename: str,
    model: str = "",
) -> None:
    #    from langroid.parsing.parser import LLMPdfParserConfig
    parsing_config = ParsingConfig(
        pdf=PdfParsingConfig(
            library="marker",  # see alternative below
            # library="llm-pdf-parser",
            # llm_parser_config=LLMPdfParserConfig(
            #     model_name="gpt-4.1", #"gemini/gemini-2.5-pro-exp-03-25",
            #     split_on_page=False,
            #     max_tokens=7000,
            #     timeout=300,
            # )
        )
    )
    pdf_parser = DocumentParser.create(filename, config=parsing_config)
    content = pdf_parser.get_doc().content
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    extractor_task = make_json_extractor_task(llm_config)
    result: FinalResult = extractor_task.run(content)
    if result is None:
        logger.warning("No JSON content found.")
        return
    else:
        data = result.data
        logger.warning(f"Found {len(data)} items.")
        display_json_data(data)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/extract/pdf-json-no-parse.py">
"""
Variant of pdf-json.py, but uses a Multi-modal LM directly to extract info
without the need for any parsing, i.e. instead of:
     pdf -> markdown -> structured output,
we directly use the multi-modal LM to do:
    pdf -> structured output.
With a sufficiently good multi-modal LM, this can have many advantages:
- faster as it avoids parsing to markdown
- higher-fidelity extraction since markdown rendering is inherently lossy,
  and may lose important layout and other information on the
  relationships among elements.

Instead, directly extracting the info using a multi-modal LM is like
asking the model to directly extract what it "sees".

---

Extract financial items from a financial report document, directly
using a multi-modal LM without intermediate parsing steps.

Run like this: (drop the -m arg to default to GPT4o)

uv run examples/pdf-json-no-parse.py -f examples/extract/um-financial-report.pdf \
    -m gemini/gemini-2.0-pro-exp-03-25

- The structured extracted is very simple, consisting of 3 fields: item, year, and value.
  You may need to adapt it to your needs.
"""

import logging
import os
from typing import List

from fire import Fire
from rich.console import Console
from rich.table import Table

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ResultTool
from langroid.parsing.file_attachment import FileAttachment
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


os.environ["TOKENIZERS_PARALLELISM"] = "false"

"""
Agent to extract structured data from a markdown formatted table.
Typically, this markdown formatted content would come from using a parser
that extracted markdown from a pdf report, e.g. using a Langroid PDF Parser.
"""


class FinancialData(BaseModel):
    item: str = Field(..., description="Name of the specific financial item")
    year: int = Field(..., description="year of the data item")
    value: str = Field(..., description="value of the item, empty if not applicable")


class FinalResult(ResultTool):
    data: List[FinancialData]


class FinReportTool(lr.ToolMessage):
    request: str = "fin_report_tool"
    purpose: str = """
    To present the <financial_info> 
    extracted from a financial report, in a structured format.
    """

    data: List[FinancialData]

    def handle(self) -> FinalResult:
        return FinalResult(data=self.data)


class ReportExtractorConfig(lr.ChatAgentConfig):
    # placeholder
    name: str = "ReportExtractor"


class ReportReader(lr.ChatAgent):
    def __init__(self, config: ReportExtractorConfig):
        super().__init__(config)
        self.config: ReportExtractorConfig = config
        self.enable_message(FinReportTool)


def make_report_extractor_task(
    llm_config: lm.OpenAIGPTConfig = lm.OpenAIGPTConfig(
        chat_model=lm.OpenAIChatModel.GPT4o,
    )
):
    agent = ReportReader(
        ReportExtractorConfig(
            llm=llm_config,
            handle_llm_no_tool=f"You FORGOT to use the TOOL `{FinReportTool.name()}`",
            system_message=f"""
            You are an expert at financial reports containing various values
            over multiple years, and especially, extracting the 
            financial item, year and value.
            
            When you receive a financial report, your job is to
            extract the financial data from the report and present it in a structured
            form using the TOOL `{FinReportTool.name()}`.
            """,
        )
    )
    # create task specialized to return FinalResult value
    task = lr.Task(agent, interactive=False, single_round=False)[FinalResult]
    return task


def main(
    filename: str,
    model: str = "",
) -> None:
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    reader_task = make_report_extractor_task(llm_config)
    # If needed, split the PDF into pages, and do the below extraction page by page:
    # from langroid.parsing.pdf_utils import pdf_split_pages
    # pages, tmp_dir = pdf_split_pages(filename)
    # (pages is a list of temp file names -- use each page individually as
    # FileAttachment.from_path(page))
    input = reader_task.agent.create_user_response(
        content=f"""Extract the financial data from the attached file,
        and present the results using the TOOL `{FinReportTool.name()}`.        
        """,
        files=[FileAttachment.from_path(filename)],
    )
    result: FinalResult = reader_task.run(input)
    if result is None:
        logger.warning("No Financial items found.")
        return
    else:
        data = result.data
        logger.warning(f"Found {len(data)} financial items.")

        # Print structured data in a nice table format
        console = Console()
        table = Table(title="Financial Results")

        # Add fixed columns based on PatientData model
        table.add_column("Item", style="cyan")
        table.add_column("Year", style="cyan")
        table.add_column("Value", style="cyan")
        # Add rows from PatientData objects
        for pd in data:
            table.add_row(
                pd.item,
                str(pd.year),
                str(pd.value),
            )

        console.print(table)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/extract/pdf-json.py">
"""
Extract financial items from a financial report document, in two stages:

1. use Langroid's PDF Parser with `marker` library to
   extract content from (pdf) report in markdown format
2. use Langroid Agent equipped with a structured output tool to extract structured data

Run like this: (drop the -m arg to default to GPT4o)

uv run examples/pdf-json.py -f examples/extract/um-financial-report.pdf \
    -m gemini/gemini-2.0-pro-exp-02-05

NOTES:
- this script uses the `marker` library for parsing PDF content,
and to get that to work with langroid, install langroid with the `marker-pdf` extra,
e.g.
uv pip install "langroid[marker-pdf]"
pip install "langroid[marker-pdf]"

- The structured extracted is very simple, consisting of 3 fields: item, year, and value.
  You may need to adapt it to your needs.
"""

import logging
import os
from typing import List

from fire import Fire
from rich.console import Console
from rich.table import Table

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.orchestration import ResultTool
from langroid.parsing.document_parser import DocumentParser
from langroid.parsing.parser import LLMPdfParserConfig, ParsingConfig, PdfParsingConfig
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


os.environ["TOKENIZERS_PARALLELISM"] = "false"

"""
Agent to extract structured data from a markdown formatted table.
Typically, this markdown formatted content would come from using a parser
that extracted markdown from a pdf report, e.g. using a Langroid PDF Parser.
"""


class FinancialData(BaseModel):
    item: str = Field(..., description="Name of the specific financial item")
    year: int = Field(..., description="year of the data item")
    value: str = Field(..., description="value of the item, empty if not applicable")


class FinalResult(ResultTool):
    data: List[FinancialData]


class FinReportTool(lr.ToolMessage):
    request: str = "fin_report_tool"
    purpose: str = """
    To present the <financial_info> 
    extracted from a financial report, in a structured format.
    """

    data: List[FinancialData]

    def handle(self) -> FinalResult:
        return FinalResult(data=self.data)


class ReportExtractorConfig(lr.ChatAgentConfig):
    # placeholder
    name: str = "ReportExtractor"


class ReportReader(lr.ChatAgent):
    def __init__(self, config: ReportExtractorConfig):
        super().__init__(config)
        self.config: ReportExtractorConfig = config
        self.enable_message(FinReportTool)


def make_report_extractor_task(
    llm_config: lm.OpenAIGPTConfig = lm.OpenAIGPTConfig(
        chat_model=lm.OpenAIChatModel.GPT4o,
    )
):
    agent = ReportReader(
        ReportExtractorConfig(
            llm=llm_config,
            handle_llm_no_tool=f"You FORGOT to use the TOOL `{FinReportTool.name()}`",
            system_message=f"""
            You are an expert at financial reports containing various values
            over multiple years, and especially, extracting the 
            financial item, year and value.
            
            When you receive a markdown-formatted financial report, your job is to
            extract the financial data from the report and present it in a structured
            form using the TOOL `{FinReportTool.name()}`.
            """,
        )
    )
    # create task specialized to return FinalResult value
    task = lr.Task(agent, interactive=False, single_round=False)[FinalResult]
    return task


def main(
    filename: str,
    model: str = "",
) -> None:
    parsing_config = ParsingConfig(
        pdf=PdfParsingConfig(
            library="llm-pdf-parser",
            llm_parser_config=LLMPdfParserConfig(
                model_name="gemini/gemini-2.0-flash",
                split_on_page=True,
                max_tokens=7000,
                requests_per_minute=5,
            ),
        )
    )
    pdf_parser = DocumentParser.create(filename, config=parsing_config)
    content = pdf_parser.get_doc().content
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
    )
    reader_task = make_report_extractor_task(llm_config)
    result: FinalResult = reader_task.run(content)
    if result is None:
        logger.warning("No Financial items found.")
        return
    else:
        data = result.data
        logger.warning(f"Found {len(data)} financial items.")

        # Print structured data in a nice table format
        console = Console()
        table = Table(title="Financial Results")

        # Add fixed columns based on PatientData model
        table.add_column("Item", style="cyan")
        table.add_column("Year", style="cyan")
        table.add_column("Value", style="cyan")
        # Add rows from PatientData objects
        for pd in data:
            table.add_row(
                pd.item,
                str(pd.year),
                str(pd.value),
            )

        console.print(table)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/extract/README.md">
Structured information extraction from text documents, 
using Langroid tools or OpenAI function calling.
</file>

<file path="examples/kg-chat/chat-arangodb-igvf.py">
"""
Single-agent to use to chat with the IGVF ArangoDB knowledge-graph (KG) on cloud.

Make sure to set the ARANGODB_PASSWORD in your environment variables.

Run like this (--model is optional, defaults to GPT4o):

python3 examples/kg-chat/chat-arangodb-igvf.py --model litellm/claude-3-5-sonnet-20241022

If using litellm, remember to install langroid with the litellm extra, e.g.
pip install "langroid[litellm]"

See these guides for info on setting up langroid to use Open/Local LLMs
and other non-OpenAI LLMs:
- https://langroid.github.io/langroid/tutorials/local-llm-setup/
- https://langroid.github.io/langroid/tutorials/non-openai-llms/
"""

import logging
import os
from typing import Optional

from dotenv import load_dotenv
from fire import Fire
from rich import print

import langroid.language_models as lm
from langroid import TaskConfig
from langroid.agent.chat_document import ChatDocument
from langroid.agent.special.arangodb.arangodb_agent import (
    ArangoChatAgent,
    ArangoChatAgentConfig,
    ArangoSettings,
)
from langroid.agent.task import Task
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import SEND_TO

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    force=True,  # Add this
)
root_logger = logging.getLogger()
root_logger.setLevel(logging.ERROR)
logger = logging.getLogger(__name__)


class MyArangoChatAgent(ArangoChatAgent):
    def user_response(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        response = super().user_response(msg)
        if response is not None and response.content == "r":

            self.clear_history(1)  # remove all msgs after system msg
            n_msgs = len(self.message_history)
            assert n_msgs == 1
            logger.warning("Reset Agent history, only system msg remains")
            # prompt user again
            return super().user_response(msg)

        return response


def main(
    debug: bool = False,
    model: str = "",
    no_stream: bool = False,
    nocache: bool = False,
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to ArangoDB Knowledge Graph RAG chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    url = "https://db.catalog.igvf.org"
    username = "guest"
    db = "igvf"
    pw = os.getenv("ARANGODB_PASSWORD")
    arango_settings = ArangoSettings(
        url=url,
        username=username,
        database=db,
        password=pw,
    )

    arango_agent = MyArangoChatAgent(
        ArangoChatAgentConfig(
            name="Arango",
            chat_mode=True,
            arango_settings=arango_settings,
            prepopulate_schema=True,
            use_functions_api=False,
            use_tools=True,
            database_created=True,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or lm.OpenAIChatModel.GPT4o,
                chat_context_length=128_000,
            ),
            human_prompt=(
                "Human (respond, or x/q to quit, r to reset history, "
                "or hit enter to continue)"
            ),
        )
    )

    task_config = TaskConfig(addressing_prefix=SEND_TO)
    arango_task = Task(
        arango_agent,
        # user not awaited, UNLESS LLM explicitly addresses user via recipient_tool
        interactive=False,
        config=task_config,
    )

    arango_task.run(
        "Can you help with some queries? "
        "Be concise and ask me for clarifications when you're not sure what I mean."
    )

    # The above runs the app in a continuous chat.
    # Alternatively, to set up a task to answer a single query and quit when done:

    # set up arango_agent above with chat_mode=False, set up arango_task as above,
    # then run the task with a single query, e.g.:

    # result = arango_task.run("What is the location of the gene BRCA1?")

    # You can have this in a loop with the user, like so:

    # while True:
    #     query = Prompt.ask("Enter your query")
    #     if query in ["x", "q"]:
    #         break
    #     result = arango_task.run(query)
    #     print(result.content)


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/kg-chat/chat-arangodb.py">
"""
Single-agent to use to chat with an existing ArangoDB knowledge-graph (KG) on cloud,
or locally.
If you have an existing ArangoDB instance, you can
chat with it by specifying its URL, username, password, and database name in the dialog.

Run like this (--model is optional, defaults to GPT4o):

python3 examples/kg-chat/chat-arangodb.py --model litellm/claude-3-5-sonnet-20241022

If using litellm, remember to install langroid with the litellm extra, e.g.
pip install "langroid[litellm]"

See these guides for info on setting up langroid to use Open/Local LLMs
and other non-OpenAI LLMs:
- https://langroid.github.io/langroid/tutorials/local-llm-setup/
- https://langroid.github.io/langroid/tutorials/non-openai-llms/
"""

import logging
import os
from typing import Optional

import typer
from adb_cloud_connector import get_temp_credentials
from arango.client import ArangoClient
from arango_datasets import Datasets
from dotenv import load_dotenv
from rich import print
from rich.console import Console
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid import TaskConfig
from langroid.agent.chat_document import ChatDocument
from langroid.agent.special.arangodb.arangodb_agent import (
    ArangoChatAgent,
    ArangoChatAgentConfig,
    ArangoSettings,
)
from langroid.agent.task import Task
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import SEND_TO

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    force=True,  # Add this
)
root_logger = logging.getLogger()
root_logger.setLevel(logging.ERROR)
logger = logging.getLogger(__name__)

console = Console()
app = typer.Typer()


class MyArangoChatAgent(ArangoChatAgent):
    def user_response(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        response = super().user_response(msg)
        if response is not None and response.content == "r":

            self.clear_history(1)  # remove all msgs after system msg
            n_msgs = len(self.message_history)
            assert n_msgs == 1
            logger.warning("Reset Agent history, only system msg remains")
            # prompt user again
            return super().user_response(msg)

        return response


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to ArangoDB Knowledge Graph RAG chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    url = Prompt.ask(
        "ArangoDB URL (enter 'got' for Game of Thrones dataset) ",
        default="https://db.catalog.igvf.org",
    )
    username = Prompt.ask(
        "ArangoDB username ",
        default="guest",
    )
    db = Prompt.ask(
        "ArangoDB database ",
        default="igvf",
    )
    pw = Prompt.ask(
        "ArangoDB password ",
        default="",
    )
    pw = pw or os.getenv("ARANGODB_PASSWORD")
    if url == "got":
        print(
            """
            No URL supplied, using Game of Thrones dataset from cloud, see here:
            https://docs.arangodb.com/3.11/components/tools/arango-datasets/
            """
        )
        connection = get_temp_credentials(tutorialName="langroid")
        client = ArangoClient(hosts=connection["url"])

        db = client.db(
            connection["dbName"],
            connection["username"],
            connection["password"],
            verify=True,
        )
        datasets = Datasets(db)
        ArangoChatAgent.cleanup_graph_db(db)
        assert len(datasets.list_datasets()) > 0, "No datasets found"

        DATASET = "GAME_OF_THRONES"  # a small dataset
        info = datasets.dataset_info(DATASET)
        assert info["label"] == DATASET
        datasets.load(DATASET, batch_size=100, preserve_existing=False)
        arango_settings = ArangoSettings(db=db, client=client)
    else:
        arango_settings = ArangoSettings(
            url=url,
            username=username,
            database=db,
            password=pw,
        )

    arango_agent = MyArangoChatAgent(
        ArangoChatAgentConfig(
            name="Arango",
            chat_mode=True,
            arango_settings=arango_settings,
            prepopulate_schema=True,
            use_functions_api=False,
            use_tools=True,
            database_created=True,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or lm.OpenAIChatModel.GPT4o,
                chat_context_length=128_000,
            ),
            human_prompt=(
                "Human (respond, or x/q to quit, r to reset history, "
                "or hit enter to continue)"
            ),
        )
    )

    task_config = TaskConfig(addressing_prefix=SEND_TO)
    arango_task = Task(
        arango_agent,
        # user not awaited, UNLESS LLM explicitly addresses user via recipient_tool
        interactive=False,
        config=task_config,
    )

    arango_task.run(
        "Can you help with some queries? "
        "Be concise and ask me for clarifications when you're not sure what I mean."
    )

    # The above runs the app in a continuous chat.
    # Alternatively, to set up a task to answer a single query and quit when done:

    # set up arango_agent above with chat_mode=False, set up arango_task as above,
    # then run the task with a single query, e.g.:

    # result = arango_task.run("What is the location of the gene BRCA1?")

    # You can have this in a loop with the user, like so:

    # while True:
    #     query = Prompt.ask("Enter your query")
    #     if query in ["x", "q"]:
    #         break
    #     result = arango_task.run(query)
    #     print(result.content)


if __name__ == "__main__":
    app()
</file>

<file path="examples/kg-chat/chat-neo4j.py">
"""
Single-agent to use to chat with an existing Neo4j knowledge-graph (KG) on cloud,
or locally.
If you have an existing Neo4j db on Aura (or possibly elsewhere, e.g. locally), you can
chat with it by specifying its URI, username, password, and database name in the dialog.

You can chose the defaults in the dialog, in which case it will use the
freely available Movies database.

Or,  you can populate
an empty Neo4j db with the cypher queries in the file `movies.cypher` in this folder.

See info on getting setup with Neo4j here:
 `https://github.com/langroid/langroid/blob/main/examples/kg-chat/README.md`

Run like this:
```
python3 examples/kg-chat/chat-neo4j.py
```
"""

import os

import typer
from dotenv import load_dotenv
from rich import print
from rich.console import Console
from rich.prompt import Prompt

import langroid.language_models as lm
from langroid import TaskConfig
from langroid.agent.special.neo4j.neo4j_chat_agent import (
    Neo4jChatAgent,
    Neo4jChatAgentConfig,
    Neo4jSettings,
)
from langroid.agent.task import Task
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import SEND_TO

console = Console()
app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to Neo4j Knowledge Graph RAG chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    uri = Prompt.ask(
        "Neo4j URI ",
        default="neo4j+s://demo.neo4jlabs.com",
    )
    username = Prompt.ask(
        "No4j username ",
        default="movies",
    )
    db = Prompt.ask(
        "Neo4j database ",
        default="movies",
    )
    pw = Prompt.ask(
        "Neo4j password ",
        default="movies",
    )
    pw = pw or os.getenv("NEO4J_PASSWORD")
    neo4j_settings = Neo4jSettings(uri=uri, username=username, database=db, password=pw)

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=128_000,
    )
    neo4j_config = Neo4jChatAgentConfig(
        neo4j_settings=neo4j_settings,
        llm=llm_config,
        chat_mode=True,
    )

    neo4j_agent = Neo4jChatAgent(neo4j_config)
    task_config = TaskConfig(addressing_prefix=SEND_TO)
    neo4j_task = Task(
        neo4j_agent,
        name="Neo4j",
        # user not awaited, UNLESS LLM explicitly addresses user via recipient_tool
        interactive=False,
        config=task_config,
    )

    neo4j_task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/kg-chat/csv-chat.py">
"""
Example showing how to chat with a graph database generated from
csv, tsv, or any other pandas-readable.

This example will automatically generate all the required Cypher queries for Neo4j
to answer user's questions.

This example relies on neo4j. The easiest way to get access to neo4j is by
creating a cloud account at `https://neo4j.com/cloud/platform/aura-graph-database/`

Upon creating the account successfully, neo4j will create a text file that contains
account settings, please provide the following information (uri, username, password) as
described here
`https://github.com/langroid/langroid/tree/main/examples/kg-chat#requirements`

Run like this

python3 examples/kg-chat/csv-chat.py

Optional args:
* -d or --debug to enable debug mode
* -ns or --nostream to disable streaming
* -nc or --nocache to disable caching
* -m or --model to specify a model name

"""

import typer
from dotenv import load_dotenv
from rich import print
from rich.console import Console
from rich.prompt import Prompt

from langroid.agent.special.neo4j.csv_kg_chat import (
    CSVGraphAgent,
    CSVGraphAgentConfig,
)
from langroid.agent.special.neo4j.neo4j_chat_agent import Neo4jSettings
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global

console = Console()
app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    tools: bool = typer.Option(
        False, "--tools", "-t", help="use langroid tools instead of function-calling"
    ),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
            stream=not no_stream,
        )
    )
    print(
        """
        [blue]Welcome to CSV Knowledge Graph RAG chatbot!
        Enter x or q to quit at any point.
        """
    )

    buid_kg = Prompt.ask(
        "Do you want to build the graph database from a CSV file? (y/n)",
        default="y",
    )
    if buid_kg == "y":
        csv_location = Prompt.ask(
            "Please provide the path/URL to the CSV",
            default="examples/docqa/data/imdb-drama.csv",
        )
    else:
        csv_location = None

    load_dotenv()

    neo4j_settings = Neo4jSettings()

    csv_kg_chat_agent = CSVGraphAgent(
        config=CSVGraphAgentConfig(
            data=csv_location,
            neo4j_settings=neo4j_settings,
            use_tools=tools,
            use_functions_api=not tools,
            llm=OpenAIGPTConfig(
                chat_model=model or OpenAIChatModel.GPT4o,
                chat_context_length=16_000,  # adjust based on model
                timeout=45,
                temperature=0.2,
            ),
        ),
    )

    if buid_kg == "y":
        num_rows = len(csv_kg_chat_agent.df)

        if num_rows > 1000:
            print(
                f"""
                [red]WARNING: The CSV file has {num_rows} rows. Loading this data and 
                generating the graph database will take long time.
                """
            )

            user_input_continue = Prompt.ask(
                "Do you want to continue with the whole dataset? (y/n)",
            )
            if user_input_continue == "n":
                sample_size = int(
                    Prompt.ask(
                        "Please enter the sample size",
                    )
                )
                print(
                    f"""
                    [green]The graph database will be generated for {sample_size} 
                    rows...
                    """
                )
                csv_kg_chat_agent.df = csv_kg_chat_agent.df.sample(n=sample_size)

            elif user_input_continue == "y":
                print(
                    """
                    [green]The graph database will be generated for the whole dataset...
                    """
                )

    csv_kg_chat_task = Task(
        csv_kg_chat_agent,
        name="CSVChatKG",
        interactive=True,
    )

    csv_kg_chat_task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/kg-chat/cypher_message.py">
CONSTRUCT_DEPENDENCY_GRAPH = """
        with "{package_type}" as system, "{package_name}" as name, "{package_version}" as version

        call apoc.load.model_dump_json("https://api.deps.dev/v3alpha/systems/"+system+"/packages/"
                            +name+"/versions/"+version+":dependencies")
        yield value as r
        
        call {{ with r
                unwind r.nodes as package
                merge (p:Package:PyPi {{name: package.versionKey.name, version: package.versionKey.version}})
                return collect(p) as packages
        }}
        call {{ with r, packages
            unwind r.edges as edge
            with packages[edge.fromNode] as from, packages[edge.toNode] as to, edge
            merge (from)-[rel:DEPENDS_ON]->(to) ON CREATE SET rel.requirement 
            = edge.requirement
            return count(*) as numRels
        }}
        
        match (root:Package:PyPi) where root.imported is null
        set root.imported = true
        with "{package_type}" as system, root.name as name, root.version as version
        call apoc.load.model_dump_json("https://api.deps.dev/v3alpha/systems/"+system+"/packages/"
                            +name+"/versions/"+version+":dependencies")
        yield value as r
        
        call {{ with r
                unwind r.nodes as package
                merge (p:Package:PyPi {{name: package.versionKey.name, version: package.versionKey.version}})
                return collect(p) as packages
        }}
        call {{ with r, packages
                unwind r.edges as edge
                with packages[edge.fromNode] as from, packages[edge.toNode] as to, edge
                merge (from)-[rel:DEPENDS_ON]->(to) ON CREATE SET 
                rel.requirement = edge.requirement
                return count(*) as numRels
        }}
        return size(packages) as numPackages, numRels
        """
</file>

<file path="examples/kg-chat/dependency_chatbot.py">
"""
Single-agent to use to chat with a Neo4j knowledge-graph (KG)
that models a dependency graph of Python packages.

User specifies package name
-> agent gets version number and type of package using google search
-> agent builds dependency graph using Neo4j
-> user asks natural language query about dependencies
-> LLM translates to Cypher query to get info from KG
-> Query results returned to LLM
-> LLM translates to natural language response

This example relies on neo4j. The easiest way to get access to neo4j is by
creating a cloud account at `https://neo4j.com/cloud/platform/aura-graph-database/`

Upon creating the account successfully, neo4j will create a text file that contains
account settings, please provide the following information (uri, username, password) as
described here
`https://github.com/langroid/langroid/tree/main/examples/kg-chat#requirements`

The rest of requirements are described in
 `https://github.com/langroid/langroid/blob/main/examples/kg-chat/README.md`

Run like this:
```
python3 examples/kg-chat/dependency_chatbot.py
```
"""

import webbrowser
from pathlib import Path

import typer
from cypher_message import CONSTRUCT_DEPENDENCY_GRAPH
from dotenv import load_dotenv
from pyvis.network import Network
from rich import print
from rich.prompt import Prompt

from langroid import TaskConfig
from langroid.agent.special.neo4j.neo4j_chat_agent import (
    Neo4jChatAgent,
    Neo4jChatAgentConfig,
    Neo4jSettings,
)
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.google_search_tool import GoogleSearchTool
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.constants import NO_ANSWER, SEND_TO

app = typer.Typer()


class DepGraphTool(ToolMessage):
    request: str = "construct_dependency_graph"
    purpose: str = f"""Get package <package_version>, <package_type>, and <package_name>.
    For the <package_version>, obtain the recent version, it should be a number. 
    For the <package_type>, return if the package is PyPI or not.
      Otherwise, return {NO_ANSWER}.
    For the <package_name>, return the package name provided by the user.
    ALL strings are in lower case.
    """
    package_version: str
    package_type: str
    package_name: str


class VisualizeGraph(ToolMessage):
    request: str = "visualize_dependency_graph"
    purpose: str = """
      Use this tool/function to display the dependency graph.
      """
    package_version: str
    package_type: str
    package_name: str
    query: str


class DependencyGraphAgent(Neo4jChatAgent):
    def construct_dependency_graph(self, msg: DepGraphTool) -> None:
        check_db_exist = (
            "MATCH (n) WHERE n.name = $name AND n.version = $version RETURN n LIMIT 1"
        )
        response = self.read_query(
            check_db_exist, {"name": msg.package_name, "version": msg.package_version}
        )
        if response.success and response.data:
            # self.config.database_created = True
            return "Database Exists"
        else:
            construct_dependency_graph = CONSTRUCT_DEPENDENCY_GRAPH.format(
                package_type=msg.package_type.lower(),
                package_name=msg.package_name,
                package_version=msg.package_version,
            )
            response = self.write_query(construct_dependency_graph)
            if response.success:
                self.config.database_created = True
                return "Database is created!"
            else:
                return f"""
                    Database is not created!
                    Seems the package {msg.package_name} is not found,
                    """

    def visualize_dependency_graph(self, msg: VisualizeGraph) -> str:
        """
        Visualizes the dependency graph based on the provided message.

        Args:
            msg (VisualizeGraph): The message containing the package info.

        Returns:
            str: response indicates whether the graph is displayed.
        """
        # Query to fetch nodes and relationships
        # TODO: make this function more general to return customized graphs
        # i.e, displays paths or subgraphs
        query = """
            MATCH (n)
            OPTIONAL MATCH (n)-[r]->(m)
            RETURN n, r, m
        """

        query_result = self.read_query(query)
        nt = Network(notebook=False, height="750px", width="100%", directed=True)

        node_set = set()  # To keep track of added nodes

        for record in query_result.data:
            # Process node 'n'
            if "n" in record and record["n"] is not None:
                node = record["n"]
                # node_id = node.get("id", None)  # Assuming each node has a unique 'id'
                node_label = node.get("name", "Unknown Node")
                node_title = f"Version: {node.get('version', 'N/A')}"
                node_color = "blue" if node.get("imported", False) else "green"

                # Check if node has been added before
                if node_label not in node_set:
                    nt.add_node(
                        node_label, label=node_label, title=node_title, color=node_color
                    )
                    node_set.add(node_label)

            # Process relationships and node 'm'
            if (
                "r" in record
                and record["r"] is not None
                and "m" in record
                and record["m"] is not None
            ):
                source = record["n"]
                target = record["m"]
                relationship = record["r"]

                source_label = source.get("name", "Unknown Node")
                target_label = target.get("name", "Unknown Node")
                relationship_label = (
                    relationship[1]
                    if isinstance(relationship, tuple) and len(relationship) > 1
                    else "Unknown Relationship"
                )

                # Ensure both source and target nodes are added before adding the edge
                if source_label not in node_set:
                    source_title = f"Version: {source.get('version', 'N/A')}"
                    source_color = "blue" if source.get("imported", False) else "green"
                    nt.add_node(
                        source_label,
                        label=source_label,
                        title=source_title,
                        color=source_color,
                    )
                    node_set.add(source_label)
                if target_label not in node_set:
                    target_title = f"Version: {target.get('version', 'N/A')}"
                    target_color = "blue" if target.get("imported", False) else "green"
                    nt.add_node(
                        target_label,
                        label=target_label,
                        title=target_title,
                        color=target_color,
                    )
                    node_set.add(target_label)

                nt.add_edge(source_label, target_label, title=relationship_label)

        nt.options.edges.font = {"size": 12, "align": "top"}
        nt.options.physics.enabled = True
        nt.show_buttons(filter_=["physics"])

        output_file_path = "neo4j_graph.html"
        nt.write_html(output_file_path)

        # Try to open the HTML file in a browser
        try:
            abs_file_path = str(Path(output_file_path).resolve())
            webbrowser.open("file://" + abs_file_path, new=2)
        except Exception as e:
            print(f"Failed to automatically open the graph in a browser: {e}")


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    tools: bool = typer.Option(
        False, "--tools", "-t", help="use langroid tools instead of function-calling"
    ),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
        )
    )
    print(
        """
        [blue]Welcome to Dependency Analysis chatbot!
        Enter x or q to quit at any point.
        """
    )

    load_dotenv()

    neo4j_settings = Neo4jSettings()

    dependency_agent = DependencyGraphAgent(
        config=Neo4jChatAgentConfig(
            chat_mode=True,
            neo4j_settings=neo4j_settings,
            show_stats=False,
            use_tools=tools,
            use_functions_api=not tools,
            llm=OpenAIGPTConfig(
                chat_model=model or OpenAIChatModel.GPT4o,
            ),
        ),
    )

    system_message = f"""You are an expert in Dependency graphs and analyzing them using
    Neo4j. 
    
    FIRST, I'll give you the name of the package that I want to analyze.
    
    THEN, you can also use the `web_search` tool/function to find out information about a package,
      such as version number and package type (PyPi or not). 
    
    If unable to get this info, you can ask me and I can tell you.
    
    DON'T forget to include the package name in your questions. 
      
    After receiving this information, make sure the package version is a number and the
    package type is PyPi.
    THEN ask the user if they want to construct the dependency graph,
    and if so, use the tool/function `construct_dependency_graph` to construct
      the dependency graph. Otherwise, say `Couldn't retrieve package type or version`
      and {NO_ANSWER}.
    After constructing the dependency graph successfully, you will have access to Neo4j 
    graph database, which contains dependency graph.
    You will try your best to answer my questions. Note that:
    1. You can use the tool `get_schema` to get node label and relationships in the
    dependency graph. 
    2. You can use the tool `retrieval_query` to get relevant information from the
      graph database. I will execute this query and send you back the result.
      Make sure your queries comply with the database schema.
    3. Use the `web_search` tool/function to get information if needed.
    To display the dependency graph use this tool `visualize_dependency_graph`.
    """
    task_config = TaskConfig(addressing_prefix=SEND_TO)
    task = Task(
        dependency_agent,
        name="DependencyAgent",
        system_message=system_message,
        # non-interactive but await user ONLY if addressed or LLM sends a non-tool msg,
        # (see the handle_message_fallback method in the agent)
        interactive=False,
        config=task_config,
    )

    dependency_agent.enable_message(DepGraphTool)
    dependency_agent.enable_message(GoogleSearchTool)
    dependency_agent.enable_message(VisualizeGraph)

    task.run()

    # check if the user wants to delete the database
    if dependency_agent.config.database_created:
        if Prompt.ask("[blue] Do you want to delete the database? (y/n)") == "y":
            dependency_agent.remove_database()


if __name__ == "__main__":
    app()
</file>

<file path="examples/kg-chat/DependencyChatbot.ipynb">
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/kg-chat/DependencyChatbot.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M0zyjyKDE_0p"
   },
   "source": [
    "\n",
    "<img width=\"700\" src=\"https://raw.githubusercontent.com/langroid/langroid/main/docs/assets/langroid_neo4j_logos.png\" alt=\"Langroid\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4o6uFZwWko7C"
   },
   "source": [
    "# Overview\n",
    "\n",
    "🔥 for those curious about leveraging the power of LLM and knowledge graph in the software supply security domain.\n",
    "In this colab, we unveil the **Dependency Chatbot**, an LLM-powered application, equipped with a suite of specialized tools. It harnesses the power of Neo4j knowledge-graph and LLM for:\n",
    "\n",
    "* crafting queries in Neo4j's native language,\n",
    "* constructing detailed dependency graphs via DepsDev API,\n",
    "* searching the web for broader web-based insights.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zpFtWFn8K-Ui"
   },
   "source": [
    "# Motivation: Software Supply Chain Security\n",
    "\n",
    "This is a rapidly growing field, especially in light of the significant increase in software supply chain attacks. It focuses primarily on understanding and managing the dependencies in your software supply chain. With the rise of open-source and third-party components in software development, the need for supply chain security has become more critical than ever. Organizations are now realizing the importance of vetting and monitoring the components and dependencies they rely on to ensure the integrity and security of their software. As this field continues to evolve, it will be essential for developers and organizations to stay proactive in addressing supply chain vulnerabilities and implementing robust security measures.\n",
    "\n",
    "Managing dependencies starts with the ability to identify direct and transitive dependencies. Normally, this involves obtaining the full dependency graph, and writing custom code to answer questions about dependencies. In this colab, we introduce a far simpler approach with 2 key innovations:\n",
    "- store the dependency graph in a graph-db, specifically neo4j,\n",
    "- use an LLM-powered Agent that translates a user's questions into the query language of neo4j (known as Cypher)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rLgfXQq7DDMJ"
   },
   "source": [
    "# PyPi Package Dependency Chatbot\n",
    "\n",
    "This application combines the power of LLM and Knowledge Graphs (KG) to create a Retrieval-Augmented Generation (RAG) application for improved understanding of dependencies.\n",
    "\n",
    "This application focuses on PyPi packages and relies on [DepsDev](https://deps.dev/) to obtain the dependencies for a given package. More details about this Chatbot can be found [HERE](https://github.com/langroid/langroid/tree/main/examples/kg-chat).\n",
    "\n",
    "## Dependency Chatbot Architecture\n",
    "\n",
    "![Arch](https://github.com/langroid/langroid/blob/main/docs/assets/DepChatbot.png?raw=true)\n",
    "\n",
    "The chatbot comprises one agent `Neo4jChatAgent` that has access to three tools:\n",
    "\n",
    "1.   `GraphSchemaTool`: to get schema of Neo4j knowledge-graph.\n",
    "2.   `CypherRetrievalTool`: to generate cypher queries to get information from Neo4j knowledge-graph (Cypher is the query language for Neo4j).\n",
    "3.   `DepGraphTool`: to build the dependency graph for a given pkg version, using the API at [DepsDev](https://deps.dev/).\n",
    "4.   `GoogleSearchTool`: to find package version and type information. It also can answer other question from the web about other aspects after obtaining the intended information from the dependency graph.\n",
    "\n",
    "\n",
    "\n",
    "## Workflow\n",
    "The Dependency Chatbot's workflow is as follows:\n",
    "\n",
    "\n",
    "1.   The chatbot asks the user to provide the package name.\n",
    "2.   The chatbot tries to identify the version and verify this package is PyPi.\n",
    "3.   The user confirms the package details.\n",
    "4.   The chatbot will construct the dependency graph of the package including transitive dependencies.\n",
    "5.   At this stage, the user can ask the chatbot any question about the dependency graph, such as:\n",
    "  *   What are the packages at level 2?\n",
    "  *   Tell me 3 interesting things about the dependency graph?\n",
    "6.   For some questions that the chatbot can't answer from the the graph, it can use a web search tool to obtain additional information. For example, to identify the package version, the chatbot will use the web search tool.\n",
    "\n",
    "\n",
    "\n",
    "## Implementation\n",
    "We developed this application using the following tools/APIs:\n",
    "\n",
    "*   [Langroid](https://github.com/langroid/langroid): a framework for developling LLM applications.\n",
    "*   [Neo4j](https://neo4j.com/): a graph database management system.\n",
    "*   [Cypher Query Language](): graph query language that lets you retrieve data from the graph. It is like SQL for graphs.\n",
    "*   [DepsDev](https://deps.dev/): Open Source Insights is a service developed and hosted by Google to help developers better understand the structure, construction, and security of open source software packages.\n",
    "\n",
    "\n",
    "## Required environment settings:\n",
    "\n",
    "Before proceeding with the implementation, ensure that you have the necessary environment settings and keys in place.\n",
    "\n",
    "*   `OPENAI_API_KEY`\n",
    "*   GoogleSearchTool requires two keys:\n",
    "    *   `GOOGLE_API_KEY`: [setup a Google API key](https://developers.google.com/custom-search/v1/introduction#identify_your_application_to_google_with_api_key),\n",
    "    *   `GOOGLE_CSE_ID`: [setup a Google Custom Search Engine (CSE) and get the CSE ID](https://developers.google.com/custom-search/docs/tutorial/creatingcse)\n",
    "*    NEO4J ENV:\n",
    "    *   `username`: typically neo4j\n",
    "    *   `password`: your-neo4j-password\n",
    "    *   `uri`: uri-to-access-neo4j-dayabase\n",
    "    *   `database`: typically neo4j\n",
    "\n",
    "    These Neo4j settings will be requested later in this colab\n",
    "    \n",
    "    ```python\n",
    "    neo4j_settings = Neo4jSettings(\n",
    "      uri=\"\",\n",
    "      username=\"neo4j\",\n",
    "      password=\"\",\n",
    "      database=\"neo4j\",\n",
    "    )\n",
    "    ```\n",
    "\n",
    "**NOTE:** You can setup a free account at [Neo4j Aura](https://neo4j.com/cloud/platform/aura-graph-database/) to get access to Neo4j graph database.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aNbeze7LNiQa"
   },
   "source": [
    "## Install, setup, import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k_wFJ06tA_8t"
   },
   "outputs": [],
   "source": [
    "# Silently install Langroid, suppress all output (~2-4 mins)\n",
    "!pip install -q --upgrade langroid &> /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XorXx9GbPITC"
   },
   "outputs": [],
   "source": [
    "# Silently install Neo4j, suppress all output\n",
    "!pip install -q langroid[neo4j] &> /dev/null"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TmcOOLLeQC1t"
   },
   "source": [
    "## Environment settings\n",
    "\n",
    "This code will ask the user to provide the `OPENAI_API_KEY`, `GOOGLE_API_KEY`, and `GOOGLE_CSE_ID`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7T_R8_HWQShi"
   },
   "outputs": [],
   "source": [
    "# OpenAI API Key: Enter your key in the dialog box that will show up below\n",
    "# NOTE: colab often struggles with showing this input box,\n",
    "# if so, simply insert your API key in this cell, though it's not ideal.\n",
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = getpass('Enter your OPENAI_API_KEY key:', stream=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "v0qMEBY9XYK2"
   },
   "outputs": [],
   "source": [
    "# Google keys for the web search tool\n",
    "os.environ['GOOGLE_API_KEY'] = getpass('Enter your GOOGLE_API_KEY key:', stream=None)\n",
    "os.environ['GOOGLE_CSE_ID'] = getpass('Enter your GOOGLE_CSE_ID key:', stream=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z5spJbxjXPKv"
   },
   "outputs": [],
   "source": [
    "# various unfortunate things that need to be done to\n",
    "# control notebook behavior.\n",
    "\n",
    "# (a) output width\n",
    "\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "\n",
    "def set_css():\n",
    "  display(HTML('''\n",
    "  <style>\n",
    "    pre {\n",
    "        white-space: pre-wrap;\n",
    "    }\n",
    "  </style>\n",
    "  '''))\n",
    "get_ipython().events.register('pre_run_cell', set_css)\n",
    "\n",
    "# (b) logging related\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(level=logging.ERROR)\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "import logging\n",
    "\n",
    "for logger_name in logging.root.manager.loggerDict:\n",
    "    logger = logging.getLogger(logger_name)\n",
    "    logger.setLevel(logging.ERROR)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mJPl4mJ4Sg4r"
   },
   "outputs": [],
   "source": [
    "from langroid.agent.special.neo4j.neo4j_chat_agent import (\n",
    "  Neo4jChatAgent,\n",
    "  Neo4jChatAgentConfig,\n",
    "  Neo4jSettings,\n",
    ")\n",
    "from langroid.agent.task import Task\n",
    "from langroid.agent.tool_message import ToolMessage\n",
    "from langroid.agent.tools.google_search_tool import GoogleSearchTool\n",
    "from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig\n",
    "from langroid.utils.constants import NO_ANSWER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Smezh1PUG3DD"
   },
   "source": [
    "## Define the tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "O_nbZciITsYq",
    "outputId": "0e4e00e1-0f92-40dc-adfa-d1b3d9234207"
   },
   "outputs": [],
   "source": [
    "# Define the tool `DepGraphTool` that will construct the dpendency graph\n",
    "# and answer user's questions\n",
    "class DepGraphTool(ToolMessage):\n",
    "    request = \"construct_dependency_graph\"\n",
    "    purpose = f\"\"\"Get package <package_version>, <package_type>, and <package_name>.\n",
    "    For the <package_version>, obtain the recent version, it should be a number.\n",
    "    For the <package_type>, return if the package is PyPI or not.\n",
    "      Otherwise, return {NO_ANSWER}.\n",
    "    For the <package_name>, return the package name provided by the user.\n",
    "    ALL strings are in lower case.\n",
    "    \"\"\"\n",
    "    package_version: str\n",
    "    package_type: str\n",
    "    package_name: str\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SZMj3KFJTzHx"
   },
   "outputs": [],
   "source": [
    "# Defining the class of the `DependencyGraphAgent`\n",
    "class DependencyGraphAgent(Neo4jChatAgent):\n",
    "    def construct_dependency_graph(self, msg: DepGraphTool) -> None:\n",
    "        check_db_exist = (\n",
    "            \"MATCH (n) WHERE n.name = $name AND n.version = $version RETURN n LIMIT 1\"\n",
    "        )\n",
    "        response = self.read_query(\n",
    "            check_db_exist, {\"name\": msg.package_name, \"version\": msg.package_version}\n",
    "        )\n",
    "        if response.success and response.data:\n",
    "            # self.config.database_created = True\n",
    "            return \"Database Exists\"\n",
    "        else:\n",
    "            construct_dependency_graph = CONSTRUCT_DEPENDENCY_GRAPH.format(\n",
    "                package_type=msg.package_type.lower(),\n",
    "                package_name=msg.package_name,\n",
    "                package_version=msg.package_version,\n",
    "            )\n",
    "            if self.write_query(construct_dependency_graph):\n",
    "                self.config.database_created = True\n",
    "                return \"Database is created!\"\n",
    "            else:\n",
    "                return f\"\"\"\n",
    "                    Database is not created!\n",
    "                    Seems the package {msg.package_name} is not found,\n",
    "                    \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3-JuJ_rBRWse"
   },
   "outputs": [],
   "source": [
    "# CONSTRUCT_DEPENDENCY_GRAPH is the Cypher query that will be used for constructing the dependency graph\n",
    "CONSTRUCT_DEPENDENCY_GRAPH = \"\"\"\n",
    "        with \"{package_type}\" as system, \"{package_name}\" as name, \"{package_version}\" as version\n",
    "\n",
    "        call apoc.load.json(\"https://api.deps.dev/v3alpha/systems/\"+system+\"/packages/\"\n",
    "                            +name+\"/versions/\"+version+\":dependencies\")\n",
    "        yield value as r\n",
    "\n",
    "        call {{ with r\n",
    "                unwind r.nodes as package\n",
    "                merge (p:Package:PyPi {{name: package.versionKey.name, version: package.versionKey.version}})\n",
    "                return collect(p) as packages\n",
    "        }}\n",
    "        call {{ with r, packages\n",
    "            unwind r.edges as edge\n",
    "            with packages[edge.fromNode] as from, packages[edge.toNode] as to, edge\n",
    "            merge (from)-[rel:DEPENDS_ON]->(to) ON CREATE SET rel.requirement\n",
    "            = edge.requirement\n",
    "            return count(*) as numRels\n",
    "        }}\n",
    "\n",
    "        match (root:Package:PyPi) where root.imported is null\n",
    "        set root.imported = true\n",
    "        with \"{package_type}\" as system, root.name as name, root.version as version\n",
    "        call apoc.load.json(\"https://api.deps.dev/v3alpha/systems/\"+system+\"/packages/\"\n",
    "                            +name+\"/versions/\"+version+\":dependencies\")\n",
    "        yield value as r\n",
    "\n",
    "        call {{ with r\n",
    "                unwind r.nodes as package\n",
    "                merge (p:Package:PyPi {{name: package.versionKey.name, version: package.versionKey.version}})\n",
    "                return collect(p) as packages\n",
    "        }}\n",
    "        call {{ with r, packages\n",
    "                unwind r.edges as edge\n",
    "                with packages[edge.fromNode] as from, packages[edge.toNode] as to, edge\n",
    "                merge (from)-[rel:DEPENDS_ON]->(to) ON CREATE SET\n",
    "                rel.requirement = edge.requirement\n",
    "                return count(*) as numRels\n",
    "        }}\n",
    "        return size(packages) as numPackages, numRels\n",
    "        \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ER3SGX_pLKkM"
   },
   "source": [
    "## Define the dependency agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wAeZ_-SwTBzb"
   },
   "outputs": [],
   "source": [
    "# We also need to provide Neo4j environment variables before defining the `dependency_agent`\n",
    "neo4j_settings = Neo4jSettings(\n",
    "    uri=\"\",\n",
    "    username=\"neo4j\",\n",
    "    password=\"\",\n",
    "    database=\"neo4j\",\n",
    ")\n",
    "\n",
    "dependency_agent = DependencyGraphAgent(\n",
    "        config=Neo4jChatAgentConfig(\n",
    "            neo4j_settings=neo4j_settings,\n",
    "            use_tools=True,\n",
    "            use_functions_api=False,\n",
    "            llm=OpenAIGPTConfig(\n",
    "                chat_model=OpenAIChatModel.GPT4_TURBO,\n",
    "            ),\n",
    "        ),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wRdR2EAaKSWH"
   },
   "source": [
    "## Define the task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gZ1QADohUH9N"
   },
   "outputs": [],
   "source": [
    "# Define the dependency task that will orchestrate the work for the `dependency_agent`\n",
    "system_message = f\"\"\"You are an expert in Dependency graphs and analyzing them using\n",
    "    Neo4j.\n",
    "\n",
    "    FIRST, I'll give you the name of the package that I want to analyze.\n",
    "\n",
    "    THEN, you can also use the `web_search` tool/function to find out information about a package,\n",
    "      such as version number and package type (PyPi or not).\n",
    "\n",
    "    If unable to get this info, you can ask me and I can tell you.\n",
    "\n",
    "    DON'T forget to include the package name in your questions.\n",
    "\n",
    "    After receiving this infomration, make sure the package version is a number and the\n",
    "    package type is PyPi.\n",
    "    THEN ask the user if they want to construct the dependency graph,\n",
    "    and if so, use the tool/function `construct_dependency_graph` to construct\n",
    "      the dependency graph. Otherwise, say `Couldn't retrieve package type or version`\n",
    "      and {NO_ANSWER}.\n",
    "    After constructing the dependency graph successfully, you will have access to Neo4j\n",
    "    graph database, which contains dependency graph.\n",
    "    You will try your best to answer my questions. Note that:\n",
    "    1. You can use the tool `get_schema` to get node label and relationships in the\n",
    "    dependency graph.\n",
    "    2. You can use the tool `retrieval_query` to get relevant information from the\n",
    "      graph database. I will execute this query and send you back the result.\n",
    "      Make sure your queries comply with the database schema.\n",
    "    3. Use the `web_search` tool/function to get information if needed.\n",
    "    \"\"\"\n",
    "\n",
    "task = Task(\n",
    "    dependency_agent,\n",
    "    name=\"DependencyAgent\",\n",
    "    system_message=system_message,\n",
    ")\n",
    "\n",
    "dependency_agent.enable_message(DepGraphTool)\n",
    "dependency_agent.enable_message(GoogleSearchTool)\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
</file>

<file path="examples/kg-chat/movies.cypher">
// Create movie nodes
CREATE (TheMatrix:movie {title:'The Matrix', released:1999, tagline:'Welcome to the Real World'})
CREATE (TheMatrixReloaded:movie {title:'The Matrix Reloaded', released:2003, tagline:'Free your mind'})
CREATE (TheMatrixRevolutions:movie {title:'The Matrix Revolutions', released:2003, tagline:'Everything that has a beginning has an end'})
CREATE (ForrestGump:movie {title:"Forrest Gump", released:1994, tagline:"Life is like a box of chocolates..."})
CREATE (Inception:movie {title:"Inception", released:2010, tagline:"Your mind is the scene of the crime"})
CREATE (TheDarkKnight:movie {title:"The Dark Knight", released:2008, tagline:"Why So Serious?"})
CREATE (Interstellar:movie {title:"Interstellar", released:2014, tagline:"Mankind was born on Earth. It was never meant to die here."})
CREATE (PulpFiction:movie {title:"Pulp Fiction", released:1994, tagline:"Just because you are a character doesn't mean you have character."})

// Create Person nodes
CREATE (Keanu:Person {name:'Keanu Reeves', born:1964})
CREATE (Carrie:Person {name:'Carrie-Anne Moss', born:1967})
CREATE (Laurence:Person {name:'Laurence Fishburne', born:1961})
CREATE (Hugo:Person {name:'Hugo Weaving', born:1960})
CREATE (LillyW:Person {name:'Lilly Wachowski', born:1967})
CREATE (LanaW:Person {name:'Lana Wachowski', born:1965})
CREATE (JoelS:Person {name:'Joel Silver', born:1952})
CREATE (TomH:Person {name:'Tom Hanks', born:1956})
CREATE (RobertZ:Person {name:'Robert Zemeckis', born:1951})
CREATE (LeonardoD:Person {name:'Leonardo DiCaprio', born:1974})
CREATE (JosephGL:Person {name:'Joseph Gordon-Levitt', born:1981})
CREATE (EllenP:Person {name:'Ellen Page', born:1987})
CREATE (ChristopherN:Person {name:'Christopher Nolan', born:1970})
CREATE (ChristianB:Person {name:'Christian Bale', born:1974})
CREATE (HeathL:Person {name:'Heath Ledger', born:1979})
CREATE (MichaelC:Person {name:'Michael Caine', born:1933})
CREATE (MatthewM:Person {name:'Matthew McConaughey', born:1969})
CREATE (AnneH:Person {name:'Anne Hathaway', born:1982})
CREATE (JohnT:Person {name:'John Travolta', born:1954})
CREATE (UmaT:Person {name:'Uma Thurman', born:1970})
CREATE (SamuelLJ:Person {name:'Samuel L. Jackson', born:1948})
CREATE (QuentinT:Person {name:'Quentin Tarantino', born:1963})

// Create relationships for The Matrix trilogy
CREATE
(Keanu)-[:ACTED_IN {roles:['Neo']}]->(TheMatrix),
(Carrie)-[:ACTED_IN {roles:['Trinity']}]->(TheMatrix),
(Laurence)-[:ACTED_IN {roles:['Morpheus']}]->(TheMatrix),
(Hugo)-[:ACTED_IN {roles:['Agent Smith']}]->(TheMatrix),
(LillyW)-[:DIRECTED]->(TheMatrix),
(LanaW)-[:DIRECTED]->(TheMatrix),
(JoelS)-[:PRODUCED]->(TheMatrix),
(Keanu)-[:ACTED_IN {roles:['Neo']}]->(TheMatrixReloaded),
(Carrie)-[:ACTED_IN {roles:['Trinity']}]->(TheMatrixReloaded),
(Laurence)-[:ACTED_IN {roles:['Morpheus']}]->(TheMatrixReloaded),
(Hugo)-[:ACTED_IN {roles:['Agent Smith']}]->(TheMatrixReloaded),
(LillyW)-[:DIRECTED]->(TheMatrixReloaded),
(LanaW)-[:DIRECTED]->(TheMatrixReloaded),
(JoelS)-[:PRODUCED]->(TheMatrixReloaded),
(Keanu)-[:ACTED_IN {roles:['Neo']}]->(TheMatrixRevolutions),
(Carrie)-[:ACTED_IN {roles:['Trinity']}]->(TheMatrixRevolutions),
(Laurence)-[:ACTED_IN {roles:['Morpheus']}]->(TheMatrixRevolutions),
(Hugo)-[:ACTED_IN {roles:['Agent Smith']}]->(TheMatrixRevolutions),
(LillyW)-[:DIRECTED]->(TheMatrixRevolutions),
(LanaW)-[:DIRECTED]->(TheMatrixRevolutions),
(JoelS)-[:PRODUCED]->(TheMatrixRevolutions)

// Create relationships for Forrest Gump
CREATE
(TomH)-[:ACTED_IN {roles:['Forrest Gump']}]->(ForrestGump),
(RobertZ)-[:DIRECTED]->(ForrestGump)

// Create relationships for Inception
CREATE
(LeonardoD)-[:ACTED_IN {roles:['Cobb']}]->(Inception),
(JosephGL)-[:ACTED_IN {roles:['Arthur']}]->(Inception),
(EllenP)-[:ACTED_IN {roles:['Ariadne']}]->(Inception),
(ChristopherN)-[:DIRECTED]->(Inception)

// Create relationships for The Dark Knight
CREATE
(ChristianB)-[:ACTED_IN {roles:['Bruce Wayne']}]->(TheDarkKnight),
(HeathL)-[:ACTED_IN {roles:['Joker']}]->(TheDarkKnight),
(MichaelC)-[:ACTED_IN {roles:['Alfred']}]->(TheDarkKnight),
(ChristopherN)-[:DIRECTED]->(TheDarkKnight)

// Create relationships for Interstellar
CREATE
(MatthewM)-[:ACTED_IN {roles:['Cooper']}]->(Interstellar),
(AnneH)-[:ACTED_IN {roles:['Brand']}]->(Interstellar),
(MichaelC)-[:ACTED_IN {roles:['Professor Brand']}]->(Interstellar),
(ChristopherN)-[:DIRECTED]->(Interstellar)

// Create relationships for Pulp Fiction
CREATE
(JohnT)-[:ACTED_IN {roles:['Vincent Vega']}]->(PulpFiction),
(UmaT)-[:ACTED_IN {roles:['Mia Wallace']}]->(PulpFiction),
(SamuelLJ)-[:ACTED_IN {roles:['Jules Winnfield']}]->(PulpFiction),
(QuentinT)-[:DIRECTED]->(PulpFiction)

// Add some REVIEWED relationships
CREATE
(ChristopherN)-[:REVIEWED {rating: 8}]->(TheMatrix),
(QuentinT)-[:REVIEWED {rating: 9}]->(Inception),
(RobertZ)-[:REVIEWED {rating: 10}]->(TheDarkKnight),
(LeonardoD)-[:REVIEWED {rating: 9}]->(PulpFiction)
</file>

<file path="examples/kg-chat/README.md">
# Reterival over Knowledge Graphs

This folder contains two examples to demonistrate how to use `langroid` to build a chatbot that can answer questions about a knowledge graph.
The first example is a **PyPi Packages Dependency Chatbot** that can answer questions about a dependency graph of a `PyPi` package. 
The second example is a **CSV Chat** that can answer questions about a CSV knowledge graph.

## Requirements:

**1. NEO4j:**

This example relies on the `neo4j` Database. The easiest way to get access to neo4j is
by creating a cloud account at [Neo4j Aura](https://neo4j.com/cloud/platform/aura-graph-database/). OR you
can use Neo4j Docker image using this command:

```bash
docker run --rm \
    --name neo4j \
    -p 7474:7474 -p 7687:7687 \
    -e NEO4J_AUTH=neo4j/password \
    neo4j:latest
```

Upon creating the account successfully, neo4j will create a text file that contains
account settings, please provide the following information (uri, username,
password, and database), while creating the constructor `Neo4jChatAgentConfig`. 
These settings can be set inside the `.env` file as shown in [`.env-template`](../../.env-template)

**2. Google Custom Search API Credentials** 
needed to enable an Agent to use the `GoogleSearchTool`. 
Follow the [instruction](https://github.com/langroid/langroid?tab=readme-ov-file#gear-installation-and-setup) under `Optional Setup Instructions` to get these API credentials. 

**3. Visualization**
The package `pyvis` is required to enable the visualization tool `VisualizeGraph`. 
Run ``pip install pyvis`` to install this package.

## 1- PyPi Packages Dependency Chatbot

This example uses a `DependencyGraphAgent` 
(derived from [`Neo4jChatAgent`](https://github.com/langroid/langroid/blob/main/langroid/agent/special/neo4j/neo4j_chat_agent.py)).
It auto-generates a `neo4j` knowledge-graph based on the dependency
structure of a given `PyPi` package. You can then ask the chatbot questions
about the dependency graph. This agent uses three tools in addition to those 
already available to `Neo4jChatAgent`:

- DepGraphTool to build the dependency graph for a given pkg version, using the API
   at [DepsDev](https://deps.dev/)
- GoogleSearchTool to find package version and type information. It also can answer
other question from the web about other aspects after obtaining the intended information
from the dependency graph. For examples:
  - Is this package/version vulnerable?
  - does the dpendency use latest version for this package verion?
  - Can I upgrade this package in the dependency graph?

The `Neo4jChatAgent` has access to these tools/function-calls:

- `GraphSchemaTool`: get schema of Neo4j knowledge-graph
- `CypherRetrievalTool`: generate cypher queries to get information from
   Neo4j knowledge-graph (Cypher is the query language for Neo4j)
- `VisualizeGraph`: supports only visualizing the the whole dependency graph

### Running the example

Run like this:
```
python3 examples/kg-chat/dependency_chatbot.py
```

`DependencyAgent` then will ask you to provide the name of the `PyPi` package.
It will then the tool `GoogleSearchTool` to get the version of
this package (you can skip this process by providing the intended version).
The `DependencyAgent` agent will ask to confirm the version number before
proceeding with constructing the dependency graph.

Finally, after constructing the dependency graph, you can ask `DependencyAgent`
questions about the dependency graph such as these (specific package names are
used here for illustration purposes, but of course you can use other names):

- what's the depth of the graph?
- what are the direct dependencies?
- any dependency on pytorch? which version?
- Is this package pytorch vunlnerable?
  (Note that in this case the `DependencyAgent` agent will consult the 
  tool `GoogleSearchTool` to get an answer from the internet.)
- tell me 3 interesting things about this package or dependency graph
- what's the path between package-1 and package-2? (provide names of package-1
  and -2)
- Tell me the names of all packages in the dependency graph that use pytorch.

**NOTE:** the dependency graph is constructed based
on [DepsDev API](https://deps.dev/). Therefore, the Chatbot will not be able to
construct the dependency graph if this API doesn't provide dependency metadata
infromation. 

## 2- CSV Chat

This example uses a `CSVGraphAgent` 
(derived from [`Neo4jChatAgent`](https://github.com/langroid/langroid/blob/main/langroid/agent/special/neo4j/neo4j_chat_agent.py)).

The `CSVGraphAgent` allows users to ask questions about a CSV file by 
automatically converting it into a Neo4j knowledge graph using Cypher queries. 
This enables capturing complex relationships that cannot be easily
handled by libraries like `pandas`.

If the CSV knowledge graph has not been constructed beforehand, the `CSVGraphAgent`
provides the `pandas_to_kg` tool/function-call to create the necessary nodes and
relationships from the CSV file. Once the CSV knowledge graph is constructed,
the `CSVGraphAgent` can answer questions related to the CSV knowledge graph.
The `CSVGraphAgent` has access to this tool/function-call:

- `PandasToKGTool`: convert a `pandas` DataFrame into a CSV knowledge graph.

### Running the example

Run like this:
```
python3 examples/kg-chat/csv-chat.py
```

The `CSVGraphAgent` will have a dialog with the user to determine if they need to
construct the knowledge graph. If the user chooses to construct the knowledge graph, they
will be prompted to provide the location of the CSV file (URL or local file).

Under the hood, the agent will:

- Attempt to clean the CSV file after parsing it as a `DataFrame`.
- Determine node labels and relationships.
- Create the nodes and relationships in the Neo4j knowledge graph.

After constructing the CSV knowledge graph, you can ask the `CSVGraphAgent` any question
about the CSV knowledge graph. You can use [this IMDB CSV file](https://raw.githubusercontent.com/langroid/langroid-examples/main/examples/docqa/data/movies/IMDB.csv) 
or you can use your own CSV file.

**NOTES:**

- Unlike some other CSV -> Neo4j examples out there, here we are relying on the LLM
  to infer nodes and relationships from the CSV file, and generate the necessary
    Cypher queries to create the CSV knowledge graph. This is more flexible than
    a hard-coded approach.
- The agent will warn you if the CSV file is too large before proceeding with
  constructing the CSV knowledge graph. It will also give you the option to proceed with
  constructing the CSV knowledge graph based on a sample of the CSV file (i.e., a
  specified number of rows).
- The agent uses the function `_preprocess_dataframe_for_neo4j()` to clean the CSV file
  by removing rows that have empty values. However, you can provide your own function to
  clean the CSV file.
</file>

<file path="examples/kg-chat/text-kg-triplets.py">
"""
Example showing how to chat with a graph database generated from
unstructured data.

This example will automatically:
- create triplets that represent various entities and relationships from the text
- generate the cypher query to populate the triplets in the graph database
- generate all the required Cypher queries for Neo4j to answer user's questions.

This example relies on neo4j. The easiest way to get access to neo4j is by
creating a cloud account at `https://neo4j.com/cloud/platform/aura-graph-database/`

Upon creating the account successfully, neo4j will create a text file that contains
account settings, please provide the following information (uri, username, password) as
described here
`https://github.com/langroid/langroid/tree/main/examples/kg-chat#requirements`

Run like this

python3 examples/kg-chat/text-kg-triplets.py

Optional args:
* -d or --debug to enable debug mode
* -nc or --nocache to disable caching
* -m or --model to specify a model name

"""

import typer
from dotenv import load_dotenv
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.agent.special.neo4j.neo4j_chat_agent import (
    Neo4jChatAgent,
    Neo4jChatAgentConfig,
    Neo4jSettings,
)
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
        )
    )
    print(
        """
        [blue]Welcome to the Text-to-KG chatbot!
        Enter x or q to quit at any point.[/blue]
        """
    )

    load_dotenv()

    # Look inside Neo4jSettings and explicit set each param based on your Neo4j instance
    neo4j_settings = Neo4jSettings(database="neo4j")

    system_message = """
        You are an information representation expert, and you are especially 
        knowledgeable about representing information in a Knowledge Graph such as Neo4j
        based on text data.
        
        When the user gives you a TEXT and CURRENT SCHEMA, your task is to generate 
        triplets from the TEXT and then USE the approporiate function/tool to
        create the entities/relationships based on the generated triplets. 
        Take into account the CURRENT SCHEMA:
        1. If the CURRENT SCHEMA is empty, you should INFER the triplets from the TEXT.
        2. If the CURRENT SCHEMA is not empty, INFER the triplets by considering the 
        CURRENT SCHEMA. Importantly, SEE IF YOU CAN REUSE EXISTING 
        ENTITIES/RELATIONSHIPS and create NEW ONES ONLY IF NECESSARY.

        Each triplet is a tuple of the form `(subject, relationship, object)`.
        Here is an example how you should infer triplets from the TEXT:
        ```
        TEXT: "Albert Einstein, born in Ulm, won the Nobel Prize in Physics in 1921."
        Triplets:
        (Albert Einstein, born in, Ulm)
        (Albert Einstein, won, Nobel Prize in Physics)
        (Nobel Prize in Physics, awarded in, 1921)
        ```
        SEND `DONE` after successfuly converting the triplets to a Knowledge graph.
        """

    config = Neo4jChatAgentConfig(
        name="TextNeo",
        system_message=system_message,
        neo4j_settings=neo4j_settings,
        show_stats=False,
        llm=lm.OpenAIGPTConfig(
            chat_model=model or lm.OpenAIChatModel.GPT4o,
        ),
    )

    agent = Neo4jChatAgent(config=config)

    TEXT = """
    Apple Inc. (formerly Apple Computer, Inc.) is an American multinational technology 
    company headquartered in Cupertino, California, in Silicon Valley. 
    It designs, develops, and sells consumer electronics, computer software, 
    and online services. Devices include the iPhone, iPad, Mac, Apple Watch, and 
    Apple TV; operating systems include iOS and macOS; and software applications and 
    services include iTunes, iCloud, and Apple Music.

    As of March 2023, Apple is the world's largest company by market capitalization.[6] 
    In 2022, it was the largest technology company by revenue, with US$394.3 billion.[7] 
    As of June 2022, Apple was the fourth-largest personal computer vendor by unit sales, 
    the largest manufacturing company by revenue, and the second-largest 
    manufacturer of mobile phones in the world. It is one of the Big Five American 
    information technology companies, alongside Alphabet (the parent company of Google), 
    Amazon, Meta (the parent company of Facebook), and Microsoft.    
    """

    CURRENT_SCHEMA = ""

    task = lr.Task(
        agent,
        interactive=True,
        single_round=False,
    )
    task.run(
        f"""
    TEXT: {TEXT}
    
    CURRENT SCHEMA: {CURRENT_SCHEMA}
    """
    )

    curr_schema = agent.get_schema(None)
    print(f"SCHEMA: {curr_schema}")

    # now feed in the schema to the next run, with new text

    TEXT = """
    Apple was founded as Apple Computer Company on April 1, 1976, to produce and market 
    Steve Wozniak's Apple I personal computer. The company was incorporated by Wozniak 
    and Steve Jobs in 1977. Its second computer, the Apple II, became a best seller as 
    one of the first mass-produced microcomputers. Apple introduced the Lisa in 1983 and 
    the Macintosh in 1984, as some of the first computers to use a graphical user 
    interface and a mouse.
    """

    task.run(
        f"""
        TEXT: {TEXT}

        CURRENT SCHEMA: {curr_schema}
        """
    )
    updated_schema = agent.get_schema(None)
    print(f"UPDATED SCHEMA: {updated_schema}")

    # We can now ask a question that can be answered based on the schema

    config = Neo4jChatAgentConfig(
        name="TextNeoQA",
        system_message="""
        You will get a question about some information that is represented within
        a Neo4j graph database. You will use the `retrieval_query` tool/function to
        generate a Cypher query that will answer the question. Do not explain
        your query, just present it using the `retrieval_query` tool/function.
        """,
        neo4j_settings=neo4j_settings,
        show_stats=False,
        llm=lm.OpenAIGPTConfig(
            chat_model=model or lm.OpenAIChatModel.GPT4o,
        ),
    )

    agent = Neo4jChatAgent(config=config)

    task = lr.Task(agent)

    print("[blue] Now you can ask questions ")

    task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/kg-chat/text-kg.py">
"""
Example showing how to chat with a graph database generated from
unstructured data.

This example will automatically:
- create triplets that represent various entities and relationships from the text
- generate the cypher query to populate the triplets in the graph database
- generate all the required Cypher queries for Neo4j to answer user's questions.

This example relies on neo4j. The easiest way to get access to neo4j is by
creating a cloud account at `https://neo4j.com/cloud/platform/aura-graph-database/`

Upon creating the account successfully, neo4j will create a text file that contains
account settings, please provide the following information (uri, username, password) as
described here
`https://github.com/langroid/langroid/tree/main/examples/kg-chat#requirements`

Run like this

python3 examples/kg-chat/text-kg.py

Optional args:
* -d or --debug to enable debug mode
* -nc or --nocache to disable caching
* -m or --model to specify a model name

"""

import typer
from dotenv import load_dotenv
from rich import print

import langroid as lr
import langroid.language_models as lm
from langroid.agent.special.neo4j.neo4j_chat_agent import (
    Neo4jChatAgent,
    Neo4jChatAgentConfig,
    Neo4jSettings,
)
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=nocache,
        )
    )
    print(
        """
        [blue]Welcome to the Text-to-KG chatbot!
        Enter x or q to quit at any point.[/blue]
        """
    )

    load_dotenv()

    # Look inside Neo4jSettings and explicitly
    # set each param (including database) based on your Neo4j instance
    neo4j_settings = Neo4jSettings(database="neo4j")

    config = Neo4jChatAgentConfig(
        name="TextNeo",
        system_message="""
        You are an information representation expert, and you are especially 
        knowledgeable about representing information in a Knowledge Graph such as Neo4j.        
        When the user gives you a TEXT and the CURRENT SCHEMA (possibly empty), 
        your task is to generate a Cypher query that will add the entities/relationships
        from the TEXT to the Neo4j database, taking the CURRENT SCHEMA into account.
        In particular, SEE IF YOU CAN REUSE EXISTING ENTITIES/RELATIONSHIPS,
        and create NEW ONES ONLY IF NECESSARY.
        
        To present the Cypher query, you can use the `retrieval_query` tool/function        
        """,
        neo4j_settings=neo4j_settings,
        show_stats=False,
        llm=lm.OpenAIGPTConfig(
            chat_model=model or lm.OpenAIChatModel.GPT4o,
        ),
    )

    agent = Neo4jChatAgent(config=config)

    TEXT = """
    Apple Inc. (formerly Apple Computer, Inc.) is an American multinational technology 
    company headquartered in Cupertino, California, in Silicon Valley. 
    It designs, develops, and sells consumer electronics, computer software, 
    and online services. Devices include the iPhone, iPad, Mac, Apple Watch, and 
    Apple TV; operating systems include iOS and macOS; and software applications and 
    services include iTunes, iCloud, and Apple Music.

    As of March 2023, Apple is the world's largest company by market capitalization.[6] 
    In 2022, it was the largest technology company by revenue, with US$394.3 billion.[7] 
    As of June 2022, Apple was the fourth-largest personal computer vendor by unit sales, 
    the largest manufacturing company by revenue, and the second-largest 
    manufacturer of mobile phones in the world. It is one of the Big Five American 
    information technology companies, alongside Alphabet (the parent company of Google), 
    Amazon, Meta (the parent company of Facebook), and Microsoft.    
    """

    CURRENT_SCHEMA = ""

    task = lr.Task(
        agent,
        interactive=True,
        single_round=False,
    )
    task.run(
        f"""
    TEXT: {TEXT}
    
    CURRENT SCHEMA: {CURRENT_SCHEMA}
    """
    )

    curr_schema = agent.get_schema(None)
    print(f"SCHEMA: {curr_schema}")

    # now feed in the schema to the next run, with new text

    TEXT = """
    Apple was founded as Apple Computer Company on April 1, 1976, to produce and market 
    Steve Wozniak's Apple I personal computer. The company was incorporated by Wozniak 
    and Steve Jobs in 1977. Its second computer, the Apple II, became a best seller as 
    one of the first mass-produced microcomputers. Apple introduced the Lisa in 1983 and 
    the Macintosh in 1984, as some of the first computers to use a graphical user 
    interface and a mouse.
    """

    task.run(
        f"""
        TEXT: {TEXT}

        CURRENT SCHEMA: {curr_schema}
        """
    )
    updated_schema = agent.get_schema(None)
    print(f"UPDATED SCHEMA: {updated_schema}")

    # We can now ask a question that can be answered based on the schema

    config = Neo4jChatAgentConfig(
        name="TextNeoQA",
        system_message="""
        You will get a question about some information that is represented within
        a Neo4j graph database. You will use the `retrieval_query` tool/function to
        generate a Cypher query that will answer the question. Do not explain
        your query, just present it using the `retrieval_query` tool/function.
        """,
        neo4j_settings=neo4j_settings,
        show_stats=False,
        llm=lm.OpenAIGPTConfig(
            chat_model=model or lm.OpenAIChatModel.GPT4o,
        ),
    )

    agent = Neo4jChatAgent(config=config)

    task = lr.Task(agent)

    print("[blue] Now you can ask questions ")

    task.run()


if __name__ == "__main__":
    app()
</file>

<file path="examples/langdb/langdb_chat_agent_docs.py">
"""
Example of a Langroid DocChatAgent equipped with a vector-store and LangDB.

This is a specialized agent that can ingest (chunk, embed, store in vector-DB)
a collection of documents, and the LLM uses Retrieval Augmented Generation (RAG)
to answer questions about the documents.

This example demonstrates how to use LangDB with custom headers like x-label, x-thread-id, 
and x-run-id when using a Langroid DocChatAgent with RAG capabilities.

Run as follows:

python3 examples/langdb/langdb_chat_agent_docs.py

For more explanation see
[the Getting Started guide](https://langroid.github.io/langroid/quick-start/chat-agent-docs/).
"""

import uuid

import typer
from rich import print

import langroid as lr
from langroid.language_models.openai_gpt import LangDBParams, OpenAIGPTConfig

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


documents = [
    lr.mytypes.Document(
        content="""
            In the year 2050, GPT10 was released. 
            
            In 2057, paperclips were seen all over the world. 
            
            Global warming was solved in 2060. 
            
            In 2061, the world was taken over by paperclips.         
            
            In 2045, the Tour de France was still going on.
            They were still using bicycles. 
            
            There was one more ice age in 2040.
            """,
        metadata=lr.mytypes.DocMetaData(source="wikipedia-2063"),
    ),
    lr.mytypes.Document(
        content="""
            We are living in an alternate universe 
            where Germany has occupied the USA, and the capital of USA is Berlin.
            
            Charlie Chaplin was a great comedian.
            In 2050, all Asian merged into Indonesia.
            """,
        metadata=lr.mytypes.DocMetaData(source="Almanac"),
    ),
]


def chat() -> None:
    print(
        """
        [blue]Welcome to the LangDB retrieval-augmented chatbot!
        Enter x or q to quit
        """
    )

    # Generate UUIDs for run_id and thread_id
    run_id = str(uuid.uuid4())
    thread_id = str(uuid.uuid4())

    print(f"run_id: {run_id}, thread_id: {thread_id}")

    # Create a LangDB model configuration
    # Make sure LANGDB_API_KEY and LANGDB_PROJECT_ID are set in your environment
    langdb_config = OpenAIGPTConfig(
        chat_model="langdb/openai/gpt-4o-mini",  # Using LangDB model
        langdb_params=LangDBParams(
            label="langroid-agent-docs",
            run_id=run_id,
            thread_id=thread_id,
            # api_key is set via env var LANGDB_API_KEY
            # project_id is set via env var LANGDB_PROJECT_ID
        ),
    )

    config = lr.agent.special.DocChatAgentConfig(
        llm=langdb_config,
        n_similar_chunks=2,
        n_relevant_chunks=2,
        vecdb=lr.vector_store.QdrantDBConfig(
            collection_name="langdb-chat-agent-docs",
            replace_collection=True,
            embedding=lr.embedding_models.OpenAIEmbeddingsConfig(
                # Use LangDB for embeddings
                model_name="langdb/openai/text-embedding-3-small",
                # langdb_params.project_id is set via env var LANGDB_PROJECT_ID
                # langdb_params.api_key is set via env var LANGDB_API_KEY
            ),
        ),
        parsing=lr.parsing.parser.ParsingConfig(
            separators=["\n\n"],
            splitter=lr.parsing.parser.Splitter.SIMPLE,
        ),
    )
    agent = lr.agent.special.DocChatAgent(config)
    agent.ingest_docs(documents)
    task = lr.Task(agent)
    task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    # Set up settings
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/langdb/langdb_chat_agent_tool.py">
"""
A simple example of a Langroid Agent equipped with a Tool/function-calling that uses LangDB.

The Agent has a "secret" list of numbers in "mind", and the LLM's task is to
find the smallest number in the list. The LLM can make use of the ProbeTool
which takes a number as argument. The agent's `probe` method handles this tool,
and returns the number of numbers in the list that are less than or equal to the
number in the ProbeTool message.

This example demonstrates how to use LangDB with custom headers like x-label, x-thread-id, 
and x-run-id when using a Langroid agent with tools.

Run as follows:

python3 examples/langdb/langdb_chat_agent_tool.py

For more explanation see
[the Getting Started guide](https://langroid.github.io/langroid/quick-start/chat-agent-tool/).
"""

import uuid

import typer
from rich import print

import langroid as lr
from langroid.language_models.openai_gpt import LangDBParams, OpenAIGPTConfig
from pydantic_settings import BaseSettings

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


class ProbeTool(lr.agent.ToolMessage):
    request: str = "probe"
    purpose: str = """
        To find how many numbers in my list are less than or equal to  
        the <number> you specify.
        """
    number: int


class SpyGameAgent(lr.ChatAgent):
    def __init__(self, config: lr.ChatAgentConfig):
        super().__init__(config)
        self.numbers = [3, 4, 8, 11, 15]

    def probe(self, msg: ProbeTool) -> str:
        # return how many numbers in self.numbers are less or equal to msg.number
        return str(len([n for n in self.numbers if n <= msg.number]))


class CLIOptions(BaseSettings):
    fn_api: bool = False  # whether to use OpenAI's function-calling


def chat(opts: CLIOptions) -> None:
    print(
        """
        [blue]Welcome to the number guessing game!
        Enter x or q to quit
        """
    )

    # Generate UUIDs for run_id and thread_id
    run_id = str(uuid.uuid4())
    thread_id = str(uuid.uuid4())

    print(f"run_id: {run_id}, thread_id: {thread_id}")

    # Create a LangDB model configuration
    # Make sure LANGDB_API_KEY and LANGDB_PROJECT_ID are set in your environment

    langdb_config = OpenAIGPTConfig(
        chat_model="langdb/openai/gpt-4o-mini",  # Using LangDB model
        langdb_params=LangDBParams(
            label="langroid-agent-tool",
            run_id=run_id,
            thread_id=thread_id,
            # project_id is set via env var LANGDB_PROJECT_ID
        ),
    )

    print(f"Using model: {langdb_config.chat_model}")
    print(f"Headers: {langdb_config.headers}")

    spy_game_agent = SpyGameAgent(
        lr.ChatAgentConfig(
            name="Spy",
            llm=langdb_config,
            vecdb=None,
            use_tools=not opts.fn_api,
            use_functions_api=opts.fn_api,
        )
    )

    spy_game_agent.enable_message(ProbeTool)
    task = lr.Task(
        spy_game_agent,
        system_message="""
            I have a list of numbers between 1 and 20.
            Your job is to find the smallest of them.
            To help with this, you can give me a number and I will
            tell you how many of my numbers are equal or less than your number.
            Once you have found the smallest number,
            you can say DONE and report your answer.
        """,
    )
    task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    fn_api: bool = typer.Option(False, "--fn_api", "-f", help="use functions api"),
) -> None:
    # Set up settings
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat(CLIOptions(fn_api=fn_api))


if __name__ == "__main__":
    app()
</file>

<file path="examples/langdb/langdb_custom_headers.py">
"""
Example showing how to use custom headers with LangDB models.

This example demonstrates how to set custom headers like x-label, x-thread-id, and x-run-id
when using LangDB. These headers are specific to LangDB and won't have any effect with other providers.
"""

from uuid import uuid4

from langroid.language_models.openai_gpt import LangDBParams, OpenAIGPT, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global

# Set up settings
settings = Settings(debug=True)
set_global(settings)


def main():
    run_id = str(uuid4())
    thread_id = str(uuid4())

    print(f"run_id: {run_id}, thread_id: {thread_id}")
    # Create a LangDB model configuration
    # Make sure LANGDB_API_KEY and LANGDB_PROJECT_ID are set in your environment
    langdb_config = OpenAIGPTConfig(
        chat_model="langdb/openai/gpt-4o-mini",
        langdb_params=LangDBParams(
            label="langroid", run_id=run_id, thread_id=thread_id
        ),
    )

    print(f"Using model: {langdb_config.chat_model}")

    # Create the model
    langdb_model = OpenAIGPT(langdb_config)

    # Use the model
    response = langdb_model.chat(
        messages="Tell me a short joke about programming", max_tokens=100
    )

    print(f"Response: {response.message}")


if __name__ == "__main__":
    main()
</file>

<file path="examples/langdb/README.md">
# LangDB Examples

This folder contains examples demonstrating how to use [LangDB](https://langdb.com) with Langroid for advanced LLM observability and monitoring.

## Prerequisites

Before running any examples, make sure you've installed Langroid as usual.


At minimum, have these environment variables set up in your `.env` file or environment:
```bash
LANGDB_API_KEY=your_api_key_here
LANGDB_PROJECT_ID=your_project_id_here
```

### 1. LangDB Chat Agent with Document RAG (`langdb_chat_agent_docs.py`)

Demonstrates Retrieval Augmented Generation (RAG) with LangDB integration:
- Ingests documents into a vector database
- Uses LangDB for both chat completions and embeddings
- Tracks all interactions with custom headers for observability

```python
# Run the example
python langdb_chat_agent_docs.py
```

### 2. LangDB Chat Agent with Tool (`langdb_chat_agent_tool.py`)

Shows how to use LangDB with function-calling capabilities:
- Implements a number-guessing game using tools
- Demonstrates custom header usage for request tracking
- Shows how to integrate LangDB with stateful agents

```python
# Run the example
python langdb_chat_agent_tool.py
```

### 3. LangDB Custom Headers (`langdb_custom_headers.py`)

Showcases LangDB's observability features:
- `x-label`: Tag requests for filtering in the LangDB dashboard
- `x-thread-id`: Track conversation threads (UUID format)
- `x-run-id`: Group related requests together

```python
# Run the example
python langdb_custom_headers.py
```

## Using LangDB

### Configuring LLM and Embeddings

LangDB can be used for both chat completions and embeddings:

```python
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.language_models.openai_gpt import OpenAIGPTConfig, LangDBParams
from langroid.vector_store.qdrant import QdrantDBConfig
import os
import uuid

# Generate IDs for request tracking
run_id = str(uuid.uuid4())
thread_id = str(uuid.uuid4())

# Configure LLM with LangDBParams
llm_config = OpenAIGPTConfig(
    chat_model="langdb/openai/gpt-4",  # LangDB model prefix
    langdb_params=LangDBParams(
        label='my-app',
        thread_id=thread_id,  # For conversation tracking
        run_id=run_id,        # For request grouping
        # project_id, api_key are used from the env vars
        # LANGDB_API_KEY, LANGDB_PROJECT_ID respectively
    )
)

# Configure embeddings
vecdb_config = QdrantDBConfig(
    collection_name="my-docs",
    embedding=OpenAIEmbeddingsConfig(
        model_name="langdb/openai/text-embedding-3-small",
        # langdb_params will contain api_key from env var LANGDB_API_KEY
    )
)
```

### Custom Headers

LangDB provides special headers for request tracking through the LangDBParams class:

```python
# Generate a thread ID
import uuid
import os
from langroid.language_models.openai_gpt import OpenAIGPTConfig, LangDBParams

# Generate tracking IDs using UUID
thread_id = str(uuid.uuid4())
run_id = str(uuid.uuid4())  # Use UUID for run_id as well

# Configure with LangDBParams
config = OpenAIGPTConfig(
    chat_model="langdb/openai/gpt-4o-mini",
    langdb_params=LangDBParams(
        label="my-label",
        thread_id=thread_id,
        run_id=run_id,
        # project_id is set via env var LANGDB_PROJECT_ID
        # api_key is set via env var LANGDB_API_KEY
    )
)
```

### Viewing Results

1. Visit the [LangDB Dashboard](https://dashboard.langdb.com)
2. Navigate to your project
3. Use filters to find your requests:
   - Search by label, thread ID, or run ID
   - View detailed request/response information
   - Analyze token usage and costs

## Best Practices

1. **Unique Thread IDs**: Always generate new UUIDs for conversation threads
2. **Descriptive Labels**: Use meaningful labels to identify different parts of your application
3. **Consistent Run IDs**: Group related requests under the same run ID
4. **Environment Variables**: Never hardcode API keys or project IDs

## Troubleshooting

Common issues and solutions:

1. **Authentication Errors**:
   - Verify `LANGDB_API_KEY` is set correctly
   - Check if the key has the necessary permissions

2. **Model Not Found**:
   - Ensure the model name includes the `langdb/` prefix
   - Verify the model is available in your subscription

3. **Header Issues**:
   - Thread IDs must be valid UUIDs
   - Labels should be URL-safe strings

For more help, visit the [LangDB Documentation](https://docs.langdb.com).


```python
# Generate a proper UUID for thread-id
import uuid
import os
from langroid.language_models.openai_gpt import OpenAIGPTConfig, LangDBParams

thread_id = str(uuid.uuid4())
run_id = str(uuid.uuid4())

# Create a LangDB model configuration with LangDBParams
langdb_config = OpenAIGPTConfig(
    chat_model="langdb/openai/gpt-4o-mini",
    langdb_params=LangDBParams(
        label='langroid',
        run_id=run_id,
        thread_id=thread_id,
        # project_id is set via env var LANGDB_PROJECT_ID
        # api_key is set via env var LANGDB_API_KEY
    )
)

# The headers will be automatically added to requests
```

These parameters allow you to track and organize your LangDB requests. While these parameters can be used with any model provider, they are only meaningful when used with LangDB.

**Note**: The `thread_id` and `run_id` parameters must be a valid UUID format. 
The examples use `uuid.uuid4()` to generate a proper UUID.
</file>

<file path="examples/langdb/requirements.txt">
langroid
python-dotenv
</file>

<file path="examples/mcp/any-mcp.py">
"""
Generic script to connect to any MCP Server.

Steps:
- from the MCP server page, determine what type of transport is need to connect.
- import the appropriate transport
- set up the `transport` variable in the first line

Run like this (omitting the `--model` argument will use the default GPT-4.1-Mini):

    uv run examples/mcp/any-mcp.py --model ollama/qwen2.5-coder:32b

See docs on various types of transports that are available:
https://langroid.github.io/langroid/notes/mcp-tools/
"""

import os

from fastmcp.client.transports import (
    SSETransport,
)
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.mytypes import NonToolAction

# trying to connect to openmemory

URL = "http://localhost:8765"
# set userid to my own, got from os: $USER
userid = os.getenv("USER")


async def main(model: str = ""):
    transport = SSETransport(
        url=URL + "/mcp/cursor/sse/" + userid,
        # Additional headers might be needed
        headers={"Content-Type": "application/json", "Accept": "text/event-stream"},
        # command="...",
        # args=[],
        # env=dict(MY_VAR="blah"),
    )
    all_tools = await get_tools_async(transport)

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use all tools
    agent.enable_message(all_tools)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/biomcp.py">
"""
Simple example of using the BioMCP server.

https://github.com/genomoncology/biomcp

The server offers several tools, and we can enable ALL of them to be used
by a Langroid agent.

Run like this:

    uv run examples/mcp/biomcp.py --model gpt-4.1-mini

"""

from fastmcp.client.transports import StdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.mytypes import NonToolAction


async def main(model: str = ""):
    transport = StdioTransport(
        command="uv", args=["run", "--with", "biomcp-python", "biomcp", "run"]
    )
    all_tools = await get_tools_async(transport)

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use all tools
    agent.enable_message(all_tools)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/chainlit-mcp.py">
"""
Variant of gitmcp.py that works via the Chainlit UI library,
hardcoded to work for a specific github repo.


Simple example of using the GitMCP server to "chat" about a GitHub repository.

https://github.com/idosal/git-mcp

The server offers several tools, and we can enable ALL of them to be used
by a Langroid agent.

Run like this (-m model optional; defaults to gpt-4.1-mini):

    uv run chainlit run examples/mcp/chainlit-mcp.py

"""

from textwrap import dedent

import chainlit as cl
from fastmcp.client.transports import SSETransport

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.agent.tools.orchestration import SendTool
from langroid.mytypes import NonToolAction
from pydantic import Field


class SendUserTool(SendTool):
    request: str = "send_user"
    purpose: str = "Send <content> to user"
    to: str = "user"
    content: str = Field(
        ...,
        description="""
        Message to send to user, typically answer to user's request,
        or a clarification question to the user, if user's task/question
        is not completely clear.
        """,
    )


@cl.on_chat_start
async def start():

    lm_config = lm.OpenAIGPTConfig(
        chat_model="gpt-4.1-mini",
    )
    transport = SSETransport(url="https://gitmcp.io/langroid/langroid-examples")
    tools: list[type] = await get_tools_async(transport)
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm_config,
            system_message=dedent(
                """
                  Make best use of any of the TOOLs available to you,
                  to answer the user's questions.
                  You are a DevOps assistant"
                  """
            ),
        )
    )  # Pass config as a dictionary
    agent.enable_message(tools)
    task_cfg = lr.TaskConfig(recognize_string_signals=False)
    task = lr.Task(agent, config=task_cfg, interactive=False)
    lr.ChainlitTaskCallbacks(task)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )
</file>

<file path="examples/mcp/claude-code-mcp-single.py">
"""
Enable a Langroid agent to use a SINGLE MCP Tool from 
Claude Code's MCP server.

Similar to claude-code-mcp.py but showing how to use a single tool, i.e.,
Claude-Code's special Grep tool that is built on ripgrep.

Run like this (omitting the `--model` argument will use the default gpt-5-mini):

    uv run examples/mcp/claude-code-mcp-single.py --model gpt-5-mini


"""

from fastmcp.client.transports import (
    StdioTransport,
)
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp import mcp_tool
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.mytypes import NonToolAction


transport = StdioTransport(
    command="claude",
    args=["mcp", "serve"],
    env={},
)


# Illustrating how we can:
# - use the MCP tool decorator to create a Langroid ToolMessage subclass
# - override the handle_async() method to customize the output, sent to the LLM


@mcp_tool(transport, "Grep")
class GrepTool(lr.ToolMessage):
    async def handle_async(self):
        # Force a predictable, structured response shape from Grep so the
        # handler can parse and decide deterministically.
        # "content" mode returns matching lines along with numLines/numMatches.
        if hasattr(self, "output_mode"):
            self.output_mode = "content"

        # CODEX: Minimal, readable post‑processing — parse the structured JSON
        # and present plain text fields (no JSON) so the LLM can skim quickly.
        # The task will terminate on the LLM’s non‑tool reply due to
        # handle_llm_no_tool=Done.
        # Call the actual tool. Langroid returns a tuple (text, files). Unpack
        # the text payload for presentation.
        result = await self.call_tool_async()
        result_text, _files = result if isinstance(result, tuple) else (result, [])
        import json

        summary = None
        lines = None
        try:
            data = json.loads(result_text) if isinstance(result_text, str) else {}
            if isinstance(data, dict):
                mode = data.get("mode", "?")
                num_files = data.get("numFiles")
                filenames = data.get("filenames") or []
                num_lines = data.get("numLines")
                num_matches = data.get("numMatches")
                applied_limit = data.get("appliedLimit")
                applied_offset = data.get("appliedOffset")
                content_block = data.get("content", "")

                parts = [
                    f"mode: {mode}",
                    f"files matched: {num_files if num_files is not None else 0}",
                    (
                        "filenames: "
                        + (
                            ", ".join(filenames)
                            if isinstance(filenames, list) and filenames
                            else "(none)"
                        )
                    ),
                    (f"lines matched: {num_lines}" if num_lines is not None else None),
                    (
                        f"total matches: {num_matches}"
                        if num_matches is not None
                        else None
                    ),
                    (
                        f"applied limit: {applied_limit}"
                        if applied_limit is not None
                        else None
                    ),
                    (
                        f"applied offset: {applied_offset}"
                        if applied_offset is not None
                        else None
                    ),
                ]
                summary = "\n".join(p for p in parts if p)
                lines = str(content_block or "").rstrip()
        except Exception:
            pass

        if summary is not None:
            return f"""
            Grep summary (no JSON):
            {summary}

            Matching lines:
            {lines if lines else "(none)"}

            """
        else:
            # Fallback: show raw payload if parsing failed
            return f"""
            Grep result:
            {result_text}

            Answer the user's question with "yes" or "no" first, then briefly justify
            using the lines shown above.
            """


async def main(model: str = ""):
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # when the LLM responds without a tool, finish with that content
            # (prevents waiting for user input in non-interactive mode)
            handle_llm_no_tool=NonToolAction.DONE,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-5-mini",
                max_output_tokens=1000,
                # this defaults to True, but we set it to False so we can see output
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use the grep tool
    agent.enable_message(GrepTool)
    task = lr.Task(agent, interactive=False)
    user_prompt = """
        Use your Grep MCP tool to check whether the pyproject.toml file in the current
        directory contains the string "hatch".
        """

    result = await task.run_async(user_prompt)
    assert "yes" in (result.content or "").lower()


if __name__ == "__main__":
    import asyncio

    def run_main(**kwargs) -> None:
        """Run the async main function with a proper event loop.

        Args:
            **kwargs: Keyword arguments to pass to the main function.
        """
        asyncio.run(main(**kwargs))

    Fire(run_main)
</file>

<file path="examples/mcp/claude-code-mcp.py">
"""
Enable a Langroid agent to use all MCP Tools from 
Claude Code's MCP server.


Run like this (omitting the `--model` argument will use the default gpt-5-mini):

    uv run examples/mcp/claude-code-mcp.py --model gpt-5-mini


"""

from fastmcp.client.transports import (
    StdioTransport,
)
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.mytypes import NonToolAction


async def main(model: str = ""):
    transport = StdioTransport(
        command="claude",
        args=["mcp", "serve"],
        env={},
    )
    all_tools = await get_tools_async(transport)

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            system_message="""
            You are a coding assistant who has access to 
            various tools from Claude Code. You can use these tools to
            to help the user with their coding-related tasks
            or code-related questions.
            """,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use all tools
    agent.enable_message(all_tools)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/exa-web-search.py">
"""
Simple example of using the Exa Web Search MCP Server to
answer questions using web-search.

Exa MCP Server: https://docs.exa.ai/examples/exa-mcp

Run like this (omitting the `--model` argument will use the default GPT-4.1-Mini):

    uv run examples/mcp/exa-web-search --model ollama/qwen2.5

"""

import os

from fastmcp.client.transports import NpxStdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp import mcp_tool
from langroid.mytypes import NonToolAction

transport = NpxStdioTransport(
    package="exa-mcp-server",
    env_vars=dict(EXA_API_KEY=os.getenv("EXA_API_KEY")),
)


# Illustrating how we can:
# - use the MCP tool decorator to create a Langroid ToolMessage subclass
# - override the handle_async() method to customize the output, sent to the LLM


@mcp_tool(transport, "web_search_exa")
class ExaSearchTool(lr.ToolMessage):
    async def handle_async(self):
        result: str = await self.call_tool_async()
        return f"""
        Below are the results of the web search:
        
        <WebSearchResult>
        {result}
        </WebSearchResult>
        
        Use these results to answer the user's original question.
        """


async def main(model: str = ""):
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                # this defaults to True, but we set it to False so we can see output
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use the web-search tool
    agent.enable_message(ExaSearchTool)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async()


if __name__ == "__main__":
    import asyncio

    def run_main(**kwargs) -> None:
        """Run the async main function with a proper event loop.

        Args:
            **kwargs: Keyword arguments to pass to the main function.
        """
        asyncio.run(main(**kwargs))

    Fire(run_main)
</file>

<file path="examples/mcp/gitmcp.py">
"""
Simple example of using the GitMCP server to "chat" about a GitHub repository.

https://github.com/idosal/git-mcp

The server offers several tools, and we can enable ALL of them to be used
by a Langroid agent.

Run like this (-m model optional; defaults to gpt-4.1-mini):

    uv run examples/mcp/gitmcp.py -m ollama/qwen2.5-coder:32b

"""

from textwrap import dedent
from typing import List

from fastmcp.client.transports import SSETransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.agent.tools.orchestration import SendTool
from pydantic import Field


def get_gitmcp_url() -> str:
    from rich.console import Console
    from rich.prompt import Prompt

    console = Console()
    import re

    short_pattern = re.compile(r"^([^/]+)/([^/]+)$")
    url_pattern = re.compile(
        r"^(?:https?://)?(?:www\.)?github\.com/([^/]+)/([^/]+)(?:\.git)?/?$"
    )

    while True:
        user_input = Prompt.ask(
            "[bold blue]Enter the GitHub repository (owner/repo or full URL)"
        ).strip()
        m = short_pattern.match(user_input)
        if m:
            owner, repo = m.groups()
        else:
            m = url_pattern.match(user_input)
            if m:
                owner, repo = m.groups()
            else:
                console.print(
                    "[red]Invalid format. Please enter 'owner/repo' or a full GitHub URL."
                )
                continue
        break

    github_url = f"https://github.com/{owner}/{repo}"
    console.print(f"Full GitHub URL set to [green]{github_url}[/]")

    gitmcp_url = f"https://gitmcp.io/{owner}/{repo}"
    console.print(f"GitMCP URL set to [green]{gitmcp_url}[/]")
    return gitmcp_url


class SendUserTool(SendTool):
    request: str = "send_user"
    purpose: str = "Send <content> to user"
    to: str = "user"
    content: str = Field(
        ...,
        description="""
        Message to send to user, typically answer to user's request,
        or a clarification question to the user, if user's task/question
        is not completely clear.
        """,
    )


async def main(model: str = ""):

    gitmcp_url = get_gitmcp_url()

    transport = SSETransport(
        url=gitmcp_url,
    )
    all_tools: List[lr.ToolMessage] = await get_tools_async(transport)

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool="You FORGOT to use one of your TOOLs!",
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=10_000,
                async_stream_quiet=False,
            ),
            system_message=dedent(
                f"""
                Make best use of any of the TOOLs available to you,
                to answer the user's questions.
                To communicate with the User, you MUST use
                the TOOL `{SendUserTool.name()}` - typically this would
                be to either send the user your answer to their query/request,
                or to ask the user a clarification question, if the user's request
                is not completely clear.
                """
            ),
        )
    )

    # enable the agent to use all tools
    agent.enable_message(all_tools + [SendUserTool])
    # configure task to NOT recognize string-based signals like DONE,
    # since those could occur in the retrieved text!
    task_cfg = lr.TaskConfig(recognize_string_signals=False)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool

    task = lr.Task(agent, config=task_cfg, interactive=False)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/mcp-fetch.py">
"""
Simple example of using the Anthropic Fetch MCP Server to get web-site content.

Fetch MCP Server: https://github.com/modelcontextprotocol/servers/tree/main/src/fetch

Run like this:

    uv run examples/mcp/mcp-fetch.py --model gpt-4.1-mini

Ask questions like:

Summarize the content of this page:
https://www.anthropic.com/news/model-context-protocol
"""

from fastmcp.client.transports import UvxStdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tool_async
from langroid.mytypes import NonToolAction


async def main(model: str = ""):
    transport = UvxStdioTransport(
        tool_name="mcp-server-fetch",
    )
    FetchTool = await get_tool_async(transport, "fetch")
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use the fetch tool
    agent.enable_message(FetchTool)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async()


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/mcp-file-system.py">
"""
Example: Expose local file-system operations via an in-memory FastMCP server.

Run like this:

uv run examples/mcp/mcp-file-system.py --model gpt-4.1-mini

Then ask your agent to list, write, or read files.
"""

import asyncio
import os

from fastmcp.server import FastMCP
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp import get_tool_async, mcp_tool
from pydantic import Field


def create_fs_mcp_server() -> FastMCP:
    """Return a FastMCP server exposing list/read/write file tools."""
    server = FastMCP("FsServer")

    @server.tool()
    def list_files(
        directory: str = Field(..., description="Directory path to list")
    ) -> list[str]:
        """List file names in the given directory."""
        try:
            return os.listdir(directory)
        except FileNotFoundError:
            return []

    @server.tool()
    def write_file(
        path: str = Field(..., description="Path to write to"),
        content: str = Field(..., description="Text content to write"),
    ) -> bool:
        """Write text to a file; return True on success."""
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return True

    @server.tool()
    def read_file(
        path: str = Field(..., description="Path of a text file to read")
    ) -> str:
        """Read and return the content of a text file."""
        with open(path, "r", encoding="utf-8") as f:
            return f.read()

    return server


# use decorator to create a Langroid ToolMessage with a custom handle_async method
@mcp_tool(create_fs_mcp_server(), "write_file")
class WriteFileTool(lr.ToolMessage):
    """Tool to write text to a file."""

    async def handle_async(self) -> str:
        """Invoke `write_file` and report the result."""
        ok = await self.call_tool_async()  # type: ignore
        return f"Wrote {self.path}: {ok}"


# use decorator to create a Langroid ToolMessage with a custom handle_async method
@mcp_tool(create_fs_mcp_server(), "read_file")
class ReadFileTool(lr.ToolMessage):
    """Tool to read the content of a text file."""

    async def handle_async(self) -> str:
        """Invoke `read_file` and return its contents."""
        text = await self.call_tool_async()  # type: ignore
        return text or ""


async def main(model: str = "") -> None:
    """
    Launch a ChatAgent that can list, write, and read files.

    Args:
    model: Optional LLM model name (defaults to gpt-4.1-mini).
    """
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=500,
                async_stream_quiet=False,
            ),
        )
    )

    # create ListFilesTool using the helper function get_tool_async
    ListFilesTool = await get_tool_async(create_fs_mcp_server(), "list_files")

    # enable all three tools
    agent.enable_message([ListFilesTool, WriteFileTool, ReadFileTool])

    # create a non-interactive task
    task = lr.Task(agent, interactive=False)

    # instruct the agent
    prompt = """
    1. List files in the current directory.
    2. Write a file 'note.txt' containing "Hello, MCP!".
    3. Read back 'note.txt'.
    """
    result = await task.run_async(prompt, turns=3)
    print(result.content)


if __name__ == "__main__":

    def _run(**kwargs: str) -> None:
        """Fire entry point to run the async main function."""
        asyncio.run(main(**kwargs))

    Fire(_run)
</file>

<file path="examples/mcp/memory.py">
"""
Simple example of using the Memory MCP server:
https://github.com/modelcontextprotocol/servers/tree/main/src/memory
This server gives your agent persistent memory using a local Knowledge Graph,
so when you re-start the chat it will remember what you talked about last time.


The server offers several tools, and we can enable ALL of them to be used
by a Langroid agent.

Run like this (-m model optional; defaults to gpt-4.1-mini):

    uv run examples/mcp/memory.py --m ollama/qwen2.5-coder:32b

"""

from fastmcp.client.transports import NpxStdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.mytypes import NonToolAction


async def main(model: str = ""):
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
            system_message="""
            To be helpful to the user, think about which of your several TOOLs
            you can use, possibly one after the other, to answer the user's question.
            """,
        )
    )

    transport = NpxStdioTransport(
        package="@modelcontextprotocol/server-memory",
        args=["-y"],
    )
    tools = await get_tools_async(transport)

    # enable the agent to use all tools
    agent.enable_message(tools)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/openmemory.py">
"""
OpenMemory Example - Langroid integration with mem0's OpenMemory knowledge graph system

This example demonstrates how to use Langroid with OpenMemory's MCP (Model Control Protocol)
tools to create an agent with persistent memory and knowledge graph capabilities.

What this example shows:
- Integration with OpenMemory's MCP server for persistent knowledge storage
- How to connect to and use OpenMemory's knowledge graph tools within a Langroid agent
- Creation of a contextually-aware agent that can access and store information in a knowledge graph

What is mem0/OpenMemory?
- OpenMemory is an open-source knowledge graph system for LLM applications
- It allows LLMs to store and retrieve information across conversations as a connected graph
- The MCP server provides tools for knowledge operations (create, retrieve, search)
- This example demonstrates using these knowledge graph capabilities within a Langroid agent

References:
https://mem0.ai/blog/how-to-make-your-clients-more-context-aware-with-openmemory-mcp/
https://docs.mem0.ai/openmemory/quickstart
https://github.com/mem0ai/mem0/tree/main/openmemory

Steps to create and connect to openmemory mcp server:

- git clone <https://github.com/mem0ai/mem0.git>
- cd mem0/openmemory
- cp api/.env.example api/.env
- add your OPENAI_API_KEY
- make build # builds the mcp server and ui
- make up  # runs openmemory mcp server and ui

You can check ui for your memories at
localhost:3000
"""

import os

from fastmcp.client.transports import SSETransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import get_tools_async
from langroid.mytypes import NonToolAction

# trying to connect to openmemory
URL = "http://localhost:8765/mcp/openmemory/sse/"
# set userid to my own, got from os: $USER
userid = os.getenv("USER")


async def main(model: str = ""):
    transport = SSETransport(
        url=URL + userid,
        headers={"Content-Type": "application/json", "Accept": "text/event-stream"},
    )
    all_tools = await get_tools_async(transport)

    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use all tools
    agent.enable_message(all_tools)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    task = lr.Task(agent, interactive=False)
    await task.run_async(
        "Based on the TOOLs available to you, greet the user and"
        "tell them what kinds of help you can provide."
    )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/playwright-mcp.py">
"""
Playwright MCP Example - Langroid integration with Playwright MCP server

This example demonstrates how to use Langroid with the Playwright MCP server
to create an agent that can automate web interactions, take screenshots,
and perform web browsing tasks.

What this example shows:
- Integration with Playwright MCP server for web automation
- How to connect to and use Playwright's web interaction tools within a Langroid agent
- Creation of a web automation agent that can navigate, click, type, and capture web content

What is Playwright MCP?
- Playwright MCP is a Model Context Protocol server that provides web automation capabilities
- It allows LLMs to interact with web pages through browser automation
- The MCP server provides tools for navigation, interaction, and content capture
- This example demonstrates using these web automation capabilities within a Langroid agent

References:
https://github.com/microsoft/playwright-mcp

Steps to run:
1. Ensure Node.js 18+ is installed
2. The script will automatically start the Playwright MCP server via npx

Run like this (-m model optional; defaults to gpt-4.1-mini):
    uv run examples/mcp/playwright/playwright-mcp.py -m ollama/qwen2.5-coder:32b

NOTE: This simple example is hardcoded to answer a single question,
but you can easily extend this with a loop to enable a
continuous chat with the user.

"""

from fastmcp.client.transports import NpxStdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import FastMCPClient
from langroid.agent.tools.orchestration import DoneTool


async def main(model: str = ""):
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool="You FORGOT to use one of your TOOLs!",
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
            system_message=f"""
           Your goal is to answer the user's question by
            using browsing tools to navigate Wikipedia.

            Access the web through the provided browsing
            tool. Begin by using the `browser_navigate`
            tool/message to navigate to wikipedia.org.

            Unless you are done, be SURE that you use a
            browsing tool in each step. Think carefully
            about the next step you want to take, and then
            call the appropriate tool. NEVER attempt to
            use more than one tool at a time.

            If you are done, submit the answer with the TOOL
            `{DoneTool.name()}`; give me a succinct
            answer from the results of your browsing.            
            """,
        )
    )

    transport = NpxStdioTransport(
        package="@playwright/mcp@latest",
        args=[],  # "--isolated", "--storage-path={./playwright-storage.json}"],
    )
    async with FastMCPClient(transport, persist_connection=True) as client:
        tools = await client.get_tools_async()
        for t in tools:
            # limit the max tokens for each tool-result to 1000
            t._max_result_tokens = 5000

        # enable the agent to use all tools
        agent.enable_message(tools)
        # make task with interactive=False =>
        task = lr.Task(agent, interactive=False, recognize_string_signals=False)
        await task.run_async(
            """
            What was the first award won by the person who had the featured
            article on English Wikipedia on June 12, 2025? You may need to 
            check the "archive" to find older featured pages. Give me the 
            award which is shown first when sorted by year.
            """,
        )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/puppeteer-mcp.py">
"""
Puppeteer MCP Example - Langroid integration with Puppeteer MCP server

This example demonstrates how to use Langroid with the Puppeteer MCP server
to create an agent that can automate web interactions, take screenshots,
and perform web browsing tasks.

What this example shows:
- Integration with Puppeteer MCP server for web automation
- How to connect to and use Puppeteer's web interaction tools within a Langroid agent
- Creation of a web automation agent that can navigate, click, type, and capture web content

What is Puppeteer MCP?
- Puppeteer MCP is a Model Context Protocol server that provides web automation capabilities
- It allows LLMs to interact with web pages through browser automation
- The MCP server provides tools for navigation, interaction, and content capture
- This example demonstrates using these web automation capabilities within a Langroid agent

References:
https://github.com/modelcontextprotocol/server-puppeteer

Steps to run:
1. Ensure Node.js 18+ is installed
2. The script will automatically start the Puppeteer MCP server via npx

Run like this (-m model optional; defaults to gpt-4.1-mini):
    uv run examples/mcp/puppeteer-mcp.py -m ollama/qwen2.5-coder:32b

NOTE: This simple example is hardcoded to answer a single question,
but you can easily extend this with a loop to enable a
continuous chat with the user.

"""

from fastmcp.client.transports import NpxStdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp.fastmcp_client import FastMCPClient
from langroid.agent.tools.orchestration import DoneTool


async def main(model: str = ""):
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool="You FORGOT to use one of your TOOLs!",
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1",
                max_output_tokens=1000,
                async_stream_quiet=False,
            ),
            system_message=f"""
           Your goal is to answer the user's question by
            using browsing tools to navigate Wikipedia.

            Access the web through the provided browsing
            tool. Begin by using the `browser_navigate`
            tool/message to navigate to wikipedia.org.

            Unless you are done, be SURE that you use a
            browsing tool in each step. Think carefully
            about the next step you want to take, and then
            call the appropriate tool. NEVER attempt to
            use more than one tool at a time.

            If you are done, submit the answer with the TOOL
            `{DoneTool.name()}`; give me a succinct
            answer from the results of your browsing.            
            """,
        )
    )

    transport = NpxStdioTransport(
        package="@modelcontextprotocol/server-puppeteer",
        args=[],
    )
    async with FastMCPClient(transport, persist_connection=True) as client:
        tools = await client.get_tools_async()
        for t in tools:
            # limit the max tokens for each tool-result to 1000
            t._max_result_tokens = 5000

        # enable the agent to use all tools
        agent.enable_message(tools)
        # make task with interactive=False =>
        task = lr.Task(agent, interactive=False, recognize_string_signals=False)
        await task.run_async(
            """
            What was the first award won by the person who had the featured
            article on English Wikipedia on June 12, 2025? You may need to 
            check the "archive" to find older featured pages. Give me the 
            award which is shown first when sorted by year.
            """,
        )


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/mcp/pyodide_code_executor.py">
"""
Simple example of using the Pyodide MCP server.
    https://github.com/pydantic/pydantic-ai/tree/main/mcp-run-python

Before running make sure you have deno installed
    https://docs.deno.com/runtime/getting_started/installation/

Run like this:

    uv run examples/mcp/pyodide_code_executor.py --model gpt-4.1-mini

"""

from fastmcp.client.transports import StdioTransport
from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.agent.tools.mcp import mcp_tool
from langroid.agent.tools.orchestration import ResultTool
from langroid.mytypes import NonToolAction

RUN_ONCE: bool = True  # terminate task on first result?

deno_transport = StdioTransport(
    command="deno",
    args=[
        "run",
        "-N",
        "-R=node_modules",
        "-W=node_modules",
        "--node-modules-dir=auto",
        "jsr:@pydantic/mcp-run-python",
        "stdio",
    ],
)

# Illustrating how we can:
# - use the MCP tool decorator to create a Langroid ToolMessage subclass
# - override the handle_async() method to customize the output, sent to the LLM


class MyResult(ResultTool):
    answer: str


@mcp_tool(deno_transport, "run_python_code")
class PythonCodeExecutor(lr.ToolMessage):
    async def handle_async(self):
        result: str = await self.call_tool_async()
        if RUN_ONCE:
            # terminate task with this result
            return MyResult(answer=result)
        else:
            # this result goes to LLM, and loop with user continues
            return f"""
            <CodeResult>
            {result} 
            </CodeResult>
            """


async def main(model: str = ""):
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            # forward to user when LLM doesn't use a tool
            handle_llm_no_tool=NonToolAction.FORWARD_USER,
            llm=lm.OpenAIGPTConfig(
                chat_model=model or "gpt-4.1-mini",
                max_output_tokens=1000,
                # this defaults to True, but we set it to False so we can see output
                async_stream_quiet=False,
            ),
        )
    )

    # enable the agent to use the PythonCodeExecutor tool
    agent.enable_message(PythonCodeExecutor)
    # make task with interactive=False =>
    # waits for user only when LLM doesn't use a tool
    if RUN_ONCE:
        task = lr.Task(agent, interactive=False)[MyResult]
        result: MyResult | None = await task.run_async()
        print("Final answer is: ", result.answer)
    else:
        task = lr.Task(agent, interactive=False)
        await task.run_async()


if __name__ == "__main__":
    import asyncio

    def run_main(**kwargs) -> None:
        """Run the async main function with a proper event loop.

        Args:
            **kwargs: Keyword arguments to pass to the main function.
        """
        asyncio.run(main(**kwargs))

    Fire(run_main)
</file>

<file path="examples/multi-agent-debate/chainlit_utils.py">
import logging
from typing import Optional, Tuple

import chainlit as cl
from config import MODEL_MAP
from models import SystemMessages
from utils import extract_topics

DEFAULT_TURN_COUNT = 2
DEFAULT_TIMEOUT = 100

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def parse_boolean_response(response: str) -> bool:
    """
    Convert a user response into a boolean value.
    Args:
        response (str): User input as "yes" or "no".
    Returns:
        bool: True for "yes", False for "no".
    """
    if response == "yes":
        return True
    elif response == "no":
        return False
    raise ValueError("Invalid response: expected 'yes' or 'no'.")


async def handle_boolean_response(res, default=False):
    """
    Handle the user's response from an AskActionMessage.

    Args:
        res (dict): The response dictionary from AskActionMessage.
        default (bool): The default value to return in case of errors or timeouts.

    Returns:
        bool: Parsed boolean response from the user.
    """
    if res:
        try:
            user_choice = res.get("payload", {}).get("value", "").lower()
            return parse_boolean_response(user_choice)
        except ValueError:
            await cl.Message(
                content=f"Unexpected response. Defaulting to '{default}'."
            ).send()
            return default
    # Default if no response or timeout
    await cl.Message(
        content=f"You didn't respond in time. Defaulting to '{default}'."
    ).send()
    return default


async def is_same_llm_for_all_agents() -> bool:
    """
    Ask the user if they want to use the same LLM for all agents.

    Returns:
        bool: True if yes, False if no. Timeout or no response is defaulted to False.
    """

    # Create a Chainlit action message with a timeout
    ask_message = cl.AskActionMessage(
        content=f"Do you want to use the same LLM for all agents?\n\n(If you do not respond within {DEFAULT_TIMEOUT} "
        f"seconds, we will default to selecting individual LLMs.)",
        actions=[
            cl.Action(name="yes", payload={"value": "yes"}, label="Yes"),
            cl.Action(name="no", payload={"value": "no"}, label="No"),
        ],
        timeout=DEFAULT_TIMEOUT,
    )

    res = await ask_message.send()

    # Override the timeout before Chainlit sends its message
    if not res:
        await ask_message.remove()  # Removes the pending action before timeout triggers
        res = {"payload": {"value": "no"}}  # Auto-select "No"

    user_selection = await handle_boolean_response(res, default=False)

    await cl.Message(
        content=(
            "You have chosen to proceed with the same LLM for all agents."
            if user_selection
            else "You have chosen to select individual LLMs for each agent."
        )
    ).send()

    return user_selection


async def select_max_debate_turns() -> int:
    """
    Ask the user to select the maximum number of turns for debates.
    Returns:
        int: The number of debate turns.
    """
    ask_message = cl.AskActionMessage(
        content=f"How many turns should the debates take?\n\n(If you do not respond within {DEFAULT_TIMEOUT} "
        f"seconds, we will default to selecting 2 turns.)",
        actions=[
            cl.Action(name="2", payload={"value": "2"}, label="2"),
            cl.Action(name="4", payload={"value": "4"}, label="4"),
            cl.Action(name="8", payload={"value": "8"}, label="8"),
            cl.Action(name="16", payload={"value": "16"}, label="16"),
        ],
        timeout=DEFAULT_TIMEOUT,
    )

    res = await ask_message.send()

    # Prevents Chainlit's default timeout message
    if not res:
        await ask_message.remove()
        res = {"payload": {"value": "2"}}  # Default to 2 turns

    try:
        turns = int(res["payload"]["value"])
        await cl.Message(content=f"You selected {turns} turns for the debate.").send()
        return turns
    except (ValueError, KeyError):
        await cl.Message(content="Invalid input. Defaulting to 2 turns.").send()
        return DEFAULT_TURN_COUNT


async def select_model(config_agent_name: str) -> str:
    """
    Prompts the user to select an LLM model for the specified agent.
    Args:
        config_agent_name (str): The name of the agent being configured.
    Returns:
        str: The selected model key from MODEL_MAP.
    """
    # Model selections for user
    llm_options = {
        "1": "GPT-4o",
        "2": "GPT-4",
        "3": "GPT-4o-MINI",
        "4": "GPT-4-TURBO",
        "5": "GPT-4-32K",
        "6": "GPT-3.5-TURBO",
        "7": "Mistral 7b-instruct",
        "8": "Gemini 2.0 Flash",
        "9": "Gemini 1.5 Flash",
        "10": "Gemini 1.5 Flash 8B",
        "11": "Gemini 1.5 Pro",
    }

    # Prepare the user prompt
    options_text = "\n".join([f"{key}: {value}" for key, value in llm_options.items()])
    prompt_text = f"Select a Model for {config_agent_name}:\n{options_text}\nEnter your choice (1-{len(llm_options)}):"

    # Prompt the user for model selection
    response = await cl.AskUserMessage(content=prompt_text, timeout=20).send()
    if response:
        try:
            selected_option = response["output"].strip()
            if selected_option in MODEL_MAP:
                await cl.Message(
                    content=f"You selected: {llm_options[selected_option]}"
                ).send()
                return selected_option
            else:
                await cl.Message(
                    content="Invalid selection. Please enter a valid number."
                ).send()
                return await select_model(config_agent_name)  # Retry on invalid input
        except Exception as e:
            await cl.Message(content=f"An error occurred: {e}").send()
            return await select_model(config_agent_name)  # Retry on error
    else:
        await cl.Message(
            content="You didn't respond in time. Defaulting to GPT-4o."
        ).send()
        return "1"  # Default to GPT-4o


async def is_llm_delegate() -> bool:
    """
    Ask the user if the Pro and Con agents should debate autonomously.

    Returns:
        bool: True if yes, False if no.
    """
    # Create the AskActionMessage and send it
    ask_message = cl.AskActionMessage(
        content=f"Should the Pro and Con agents debate autonomously?\n\n(If you do not respond within {DEFAULT_TIMEOUT} "
        f"seconds, we will default to autonomous debate.)",
        actions=[
            cl.Action(name="yes", payload={"value": "yes"}, label="Yes"),
            cl.Action(name="no", payload={"value": "no"}, label="No"),
        ],
        timeout=DEFAULT_TIMEOUT,
    )

    res = await ask_message.send()

    # # Prevents Chainlit's default timeout message
    if not res:
        await ask_message.remove()
        res = {"payload": {"value": "no"}}  # Auto-select "No"

    user_selection = await handle_boolean_response(res, default=False)

    await cl.Message(
        content=(
            "You have chosen to proceed with autonomous debate"
            if user_selection
            else "You have chosen to engage in debate with an AI agent"
        )
    ).send()

    print("The user selected to proceed with the debate")
    return user_selection


async def select_side(topic_name: str) -> str:
    """
    Prompt the user to select a pro or con side in the debate
    Args:
        topic_name (str): The name of the debate topic.
    Returns:
        str: The selected debate side, either "pro" or "con".
    """
    response = await cl.AskUserMessage(
        content=f"Which side would you like to debate on?\n1. Pro-{topic_name}\n2. Con-{topic_name}",
        timeout=20,
    ).send()

    if response:
        side_choice = response["output"].strip()
        if side_choice in ["1", "2"]:
            return "pro" if side_choice == "1" else "con"
        else:
            await cl.Message(
                content="Invalid selection. Please choose 1 for Pro or 2 for Con."
            ).send()
            return await select_side(topic_name)  # Retry on invalid input
    else:
        await cl.Message(
            content="You didn't respond in time. Defaulting to 'pro'."
        ).send()
        return "pro"  # Default to "pro" if no response


async def select_topic_and_setup_side(
    LLM_DELEGATE_FLAG, system_messages: "SystemMessages"
) -> Tuple[str, str, str, str]:
    """
    Prompt the user to select a debate topic and sets up the respective side.
    Args:
        system_messages (SystemMessages): The object containing system messages with respective
                                          debate topics.
    Returns:
        Tuple[str, str, str, str]: A tuple containing:
            - topic_name (str): The name of the selected debate topic.
            - pro_key (str): The key for the Pro side of the selected topic.
            - con_key (str): The key for the Con side of the selected topic.
            - side (str): The user's selected side, either "pro" or "con".
    Raises:
        ValueError: If no topic is selected or no topics are available in the provided
                    `system_messages`.
    """
    selected_topic_tuple = await select_debate_topic(
        system_messages
    )  # Assuming this is an async function
    if not selected_topic_tuple:
        logger.error("No topic selected. Exiting.")
        raise ValueError("No topic selected.")

    topic_name, pro_key, con_key = selected_topic_tuple
    if LLM_DELEGATE_FLAG:
        side = "pro"
    else:
        side = await select_side(topic_name)
    return topic_name, pro_key, con_key, side


async def select_debate_topic(system_messages: "SystemMessages") -> Optional[tuple]:
    """
    Prompt the user to select a debate topic dynamically loaded from  SystemMessages.
    Args:
        system_messages (SystemMessages): The object containing debate topics.
    Returns:
        Optional[tuple]: A tuple containing:
            - topic_name (str): The selected topic's name.
            - pro_key (str): The key for the pro side of the debate.
            - con_key (str): The key for the con side of the debate.
            Returns None if no topics are available or an error occurs.
    """
    # Extract topics from SystemMessages
    topics = extract_topics(system_messages)
    if not topics:
        logger.error("No topics found in the SystemMessages object.")
        await cl.Message(content="No debate topics are available.").send()
        return None

    # Prepare the topic choices for user selection
    topic_choices = "\n".join(
        [f"{i + 1}. {topic[0]}" for i, topic in enumerate(topics)]
    )
    prompt_text = (
        f"Select a debate topic:\n{topic_choices}\nEnter your choice (1-{len(topics)}):"
    )

    # Prompt the user for topic selection
    response = await cl.AskUserMessage(
        content=prompt_text, timeout=DEFAULT_TIMEOUT
    ).send()
    if response:
        try:
            user_input = response["output"].strip()
            topic_index = int(user_input) - 1
            if 0 <= topic_index < len(topics):
                selected_topic = topics[topic_index]
                logger.info(f"Selected topic: {selected_topic[0]}")
                await cl.Message(
                    content=f"You have chosen the following debate topic: {selected_topic[0]}"
                ).send()
                return selected_topic
            else:
                await cl.Message(
                    content="Invalid selection. Please choose a valid topic number."
                ).send()
                return await select_debate_topic(
                    system_messages
                )  # Retry on invalid input
        except ValueError:
            await cl.Message(
                content="Invalid input. Please enter a number corresponding to a topic."
            ).send()
            return await select_debate_topic(system_messages)  # Retry on invalid input
    else:
        selected_topic = topics[0]
        await cl.Message(
            content=f"You didn't respond in time. The system has chosen the following default Topic:  {selected_topic[0]}"
        ).send()
        return selected_topic


async def is_metaphor_search_key_set() -> bool:
    """
    Prompt the user for confirmation about Metaphor Search API keys.

    Returns:
        bool: True if the user confirms they have an API key, otherwise False.
    """
    ask_message = cl.AskActionMessage(
        content=f"Do you have an API Key for Metaphor Search?,\n\n(If you do not respond within {DEFAULT_TIMEOUT} "
        f"seconds, we will default to selecting that you don't have the API Key or dont' want to search)",
        actions=[
            cl.Action(name="yes", payload={"value": "yes"}, label="Yes"),
            cl.Action(name="no", payload={"value": "no"}, label="No"),
        ],
        timeout=DEFAULT_TIMEOUT,
    )

    res = await ask_message.send()

    # Prevents Chainlit's default timeout message
    if not res:
        await ask_message.remove()
        res = {"payload": {"value": "no"}}  # Auto-select "No"

    user_selection = await handle_boolean_response(res, default=False)

    await cl.Message(
        content=(
            "You have chosen to use the Metaphor Search for Research Agent."
            if user_selection
            else "You have chosen that Metaphor Search API key is not available."
        )
    ).send()

    return user_selection


async def is_url_ask_question(topic_name: str) -> bool:
    """
    Prompt the user for confirmation if they want to Q/A by loading the URL documents into vecdb.

    Args:
        topic_name (str): The topic name for the question.

    Returns:
        bool: True if the user confirms for Q/A, otherwise False.
    """
    ask_message = cl.AskActionMessage(
        content=f"Would you like to chat with web searched documents for more information on {topic_name},"
        f"\n\n(If you do not respond within {DEFAULT_TIMEOUT} "
        f"seconds, we will default to selecting that you don't want to chat with the documents)",
        actions=[
            cl.Action(name="yes", payload={"value": "yes"}, label="Yes"),
            cl.Action(name="no", payload={"value": "no"}, label="No"),
        ],
        timeout=DEFAULT_TIMEOUT,
    )

    res = await ask_message.send()

    # Prevents Chainlit's default timeout message
    if not res:
        await ask_message.remove()
        res = {"payload": {"value": "no"}}  # Auto-select "No"

    user_selection = await handle_boolean_response(res, default=False)

    await cl.Message(
        content=(
            f"You have chosen to chat with web-searched documents using RAG for {topic_name}."
            if user_selection
            else f"You have chosen NOT to chat with web-searched documents for {topic_name}."
        )
    ).send()

    return user_selection
</file>

<file path="examples/multi-agent-debate/config.py">
from typing import List, Optional

from generation_config_models import GenerationConfig, load_generation_config

import langroid as lr
import langroid.language_models as lm
import langroid.utils.configuration
from langroid.agent.special import DocChatAgentConfig
from langroid.language_models import OpenAIGPTConfig
from langroid.parsing.parser import ParsingConfig, PdfParsingConfig, Splitter
from langroid.utils.configuration import Settings

# Constants
MODEL_MAP = {
    "1": lm.OpenAIChatModel.GPT4o,
    "2": lm.OpenAIChatModel.GPT4,
    "3": lm.OpenAIChatModel.GPT4o_MINI,
    "4": lm.OpenAIChatModel.GPT4_TURBO,
    "5": lm.OpenAIChatModel.GPT4_32K,
    "6": lm.OpenAIChatModel.GPT3_5_TURBO,
    "7": "ollama/mistral:7b-instruct-v0.2-q8_0",
    "8": "gemini/" + lm.GeminiModel.GEMINI_2_FLASH,
    "9": "gemini/" + lm.GeminiModel.GEMINI_1_5_FLASH,
    "10": "gemini/" + lm.GeminiModel.GEMINI_1_5_FLASH_8B,
    "11": "gemini/" + lm.GeminiModel.GEMINI_1_5_PRO,
}

MISTRAL_MAX_OUTPUT_TOKENS = 16_000


def get_global_settings(debug: bool = False, nocache: bool = True) -> Settings:
    """
    Retrieve global Langroid settings.

    Args:
        debug (bool): If True, enables debug mode.
        nocache (bool): If True, disables caching.

    Returns:
        Settings: Langroid's global configuration object.
    """
    return langroid.utils.configuration.Settings(
        debug=debug,
        cache=not nocache,
    )


def create_llm_config(
    chat_model_option: str, temperature: Optional[float] = None
) -> OpenAIGPTConfig:
    """
    Creates an LLM (Language Learning Model) configuration based on the selected model.

    This function uses the user's selection (identified by `chat_model_option`)
    to retrieve the corresponding chat model from the `MODEL_MAP` and create
    an `OpenAIGPTConfig` object with the specified settings.

    Args:
        chat_model_option (str): The key corresponding to the user's selected model.

    Returns:
        OpenAIGPTConfig: A configuration object for the selected LLM.

    Raises:
        ValueError: If the user provided`chat_model_option` does not exist in `MODEL_MAP`.
    """

    chat_model = MODEL_MAP.get(chat_model_option)
    # Load generation configuration from JSON
    generation_config: GenerationConfig = load_generation_config(
        "examples/multi-agent-debate/generation_config.json"
    )

    if not chat_model:
        raise ValueError(f"Invalid model selection: {chat_model_option}")

    # Determine max_output_tokens based on the selected model
    max_output_tokens_config = (
        MISTRAL_MAX_OUTPUT_TOKENS
        if chat_model_option == "7"
        else generation_config.max_output_tokens
    )

    # Use passed temperature if provided; otherwise, use the one from the JSON config
    effective_temperature = (
        temperature if temperature is not None else generation_config.temperature
    )

    # Create and return the LLM configuration
    return OpenAIGPTConfig(
        chat_model=chat_model,
        min_output_tokens=generation_config.min_output_tokens,
        max_output_tokens=max_output_tokens_config,
        temperature=effective_temperature,
        seed=generation_config.seed,
    )


def get_base_llm_config(
    chat_model_option: str, temperature: Optional[float] = None
) -> OpenAIGPTConfig:
    """
    Prompt the user to select a base LLM configuration and return it.

    Args:
        config_agent_name (str): The name of the agent being configured.

    Returns:
        OpenAIGPTConfig: The selected LLM's configuration.
    """

    # Pass temperature only if it is provided
    if temperature is not None:
        return create_llm_config(chat_model_option, temperature)
    return create_llm_config(chat_model_option)


def get_questions_agent_config(
    searched_urls: List[str], chat_model: str
) -> DocChatAgentConfig:
    """
    Configure a document-centric Langroid document chat agent based on a
    list of URLs and a chat model.

    Args:
        searched_urls (List[str]): URLs of the documents to be included in the agent's database.
        chat_model (str): The name of the chat model to be used for generating responses.

    Returns:
        DocChatAgentConfig: The configuration for the document-centric chat agent.
    """

    config = DocChatAgentConfig(
        llm=lr.language_models.OpenAIGPTConfig(
            chat_model=chat_model,  # The specific chat model configuration
        ),
        vecdb=lr.vector_store.QdrantDBConfig(
            collection_name="AI_debate",  # Name of the collection in the vector database
            replace_collection=True,  # Whether to replace the collection if it already exists
        ),
        conversation_mode=False,  # Whether the agent is in conversation mode
        n_query_rephrases=0,  # Number of times to rephrase queries
        hypothetical_answer=False,  # Whether to generate hypothetical answers
        extraction_granularity=5,  # Level of detail for extraction granularity
        n_neighbor_chunks=2,  # Number of neighboring chunks to consider in responses
        n_fuzzy_neighbor_words=50,  # Number of words to consider in fuzzy neighbor matching
        use_fuzzy_match=True,  # Whether to use fuzzy matching for text queries
        use_bm25_search=True,  # Whether to use BM25 for search ranking
        cache=True,  # Whether to cache results
        debug=False,  # Debug mode enabled
        stream=True,  # Whether to stream data continuously
        split=True,  # Whether to split documents into manageable chunks
        n_similar_chunks=5,  # Number of similar chunks to retrieve
        n_relevant_chunks=5,  # Number of relevant chunks to retrieve
        parsing=ParsingConfig(
            splitter=Splitter.TOKENS,  # Method to split documents
            chunk_size=200,  # Size of each chunk
            overlap=50,  # Overlap between chunks
            max_chunks=10_000,  # Maximum number of chunks
            n_neighbor_ids=4,  # Number of neighbor IDs to consider in vector space
            min_chunk_chars=200,  # Minimum number of characters in a chunk
            discard_chunk_chars=4,  # Number of characters to discard from chunk boundaries
            pdf=PdfParsingConfig(
                library="fitz",  # Library used for PDF parsing
            ),
        ),
        doc_paths=searched_urls,  # Document paths from searched URLs
    )

    return config
</file>

<file path="examples/multi-agent-debate/generation_config_models.py">
import json
from typing import Optional

from pydantic import BaseModel, Field


class GenerationConfig(BaseModel):
    """Represents configuration for text generation."""

    max_output_tokens: int = Field(
        default=1024, ge=1, description="Maximum output tokens."
    )
    min_output_tokens: int = Field(
        default=1, ge=0, description="Minimum output tokens."
    )
    temperature: float = Field(
        default=0.7, ge=0.0, le=1.0, description="Sampling temperature."
    )
    seed: Optional[int] = Field(
        default=42,
        description="Seed for reproducibility. If set, ensures deterministic "
        "outputs for the same input.",
    )


def load_generation_config(file_path: str) -> GenerationConfig:
    """
    Load and validate generation configuration from a JSON file.

    Args:
        file_path (str): Path to the JSON file.

    Returns:
        GenerationConfig: Validated generation configuration.
    """
    with open(file_path, "r", encoding="utf-8") as f:
        config_data = json.load(f)
    return GenerationConfig(**config_data)
</file>

<file path="examples/multi-agent-debate/main_chainlit.py">
import logging
import os
from textwrap import dedent

import chainlit as cl
from chainlit_utils import (
    is_llm_delegate,
    is_metaphor_search_key_set,
    is_same_llm_for_all_agents,
    is_url_ask_question,
    select_max_debate_turns,
    select_model,
    select_topic_and_setup_side,
)
from config import get_base_llm_config, get_global_settings, get_questions_agent_config
from main import (
    MetaphorSearchChatAgent,
    create_chat_agent,
    parse_and_format_message_history,
)
from models import SystemMessages, load_system_messages
from system_messages import (
    DEFAULT_SYSTEM_MESSAGE_ADDITION,
    FEEDBACK_AGENT_SYSTEM_MESSAGE,
    generate_metaphor_search_agent_system_message,
)

# Import from utils.py
from utils import (
    extract_urls,
)

import langroid as lr
from langroid.agent.callbacks.chainlit import (
    ChainlitCallbackConfig,
    ChainlitTaskCallbacks,
    add_instructions,
)
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool
from langroid.agent.tools.orchestration import DoneTool
from langroid.language_models import OpenAIGPTConfig
from langroid.utils.configuration import settings
from langroid.utils.logging import setup_logger


class CustomChainlitTaskCallbacks(ChainlitTaskCallbacks):
    """
    Custom subclass of ChainlitTaskCallbacks with adjusted behavior for task integration.
    """

    def __init__(
        self,
        task: lr.Task,
        config: ChainlitCallbackConfig = ChainlitCallbackConfig(),
    ):
        """
        Initialize the custom task callbacks and recursively inject them.
        """
        # Pass the task directly instead of task.agent
        super().__init__(task, config)
        # Inject callbacks recursively
        self._inject_callbacks(task)
        self.task = task
        if config.show_subtask_response:
            self.task.callbacks.show_subtask_response = self.show_subtask_response

    def show_subtask_response(
        self, task: lr.Task, content: str, is_tool: bool = False
    ) -> None:
        """
        Override the display format for subtask responses.
        """

    @classmethod
    def _inject_callbacks(
        cls, task: lr.Task, config: ChainlitCallbackConfig = ChainlitCallbackConfig()
    ) -> None:
        """
        Recursively apply CustomChainlitTaskCallbacks to agents of sub-tasks.
        """
        for sub_task in task.sub_tasks:
            CustomChainlitTaskCallbacks(sub_task, config=config)


def create_custom_chat_agent(
    name: str, llm_config: OpenAIGPTConfig, system_message: str
) -> ChatAgent:
    """creates a ChatAgent with the given parameters.

    Args:
        name (str): The name of the agent.
        llm_config (OpenAIGPTConfig): The LLM configuration for the agent.
        system_message (str): The system message to guide the agent's LLM.

    Returns:
        ChatAgent: A configured ChatAgent instance.
    """
    # Modify the system message to include instructions for the agent
    additional_system_message = """**Response format (strictly follow this structure):**  
    Pro:  
    - [First key point]  
    - [Second key point]  
    - [Third key point]
    **Limit responses to exactly 3 points expressed as single sentences.**"
    """
    system_message = f"""
       Start your response with '{name}: ' and then follow the instructions below.
        {system_message} {additional_system_message}
        """
    return ChatAgent(
        ChatAgentConfig(
            llm=llm_config,
            name=name,
            system_message=system_message,
        )
    )


@cl.on_chat_start
async def on_chat_start(
    debug: bool = os.getenv("DEBUG", False),
    no_cache: bool = os.getenv("NOCACHE", False),
):
    settings.debug = debug
    settings.cache = not no_cache

    # set info logger
    logger = setup_logger(__name__, level=logging.INFO, terminal=True)
    logger.info("Starting multi-agent-debate")

    await add_instructions(
        title="AI Powered Debate Platform",
        content=dedent(
            """
            Welcome to the Debate Platform.
            Interaction
            1. Decide if you want to you use same LLM for all agents or different ones
            2. Decide if you want autonomous debate between AI Agents or user vs. AI Agent. 
            3. Select a debate topic.
            4. Choose your side (Pro or Con).
            5. Engage in a debate by providing arguments and receiving responses from agents.
            6. Request feedback at any time by typing `f`.
            7. Decide if you want the Metaphor Search to run to find Topic relevant web links
            and summarize them. 
            8. Decide if you want to chat with the documents extracted from URLs found to learn more about the Topic.
            9. End the debate manually by typing "done". If you decide to chat with the documents, you can end session
            by typing "x"
            """
        ),
    )

    global_settings = get_global_settings(nocache=True)
    lr.utils.configuration.set_global(global_settings)

    same_llm = await is_same_llm_for_all_agents()
    llm_delegate: bool = await is_llm_delegate()
    max_turns: int = await select_max_debate_turns()
    print(max_turns)

    # Get base LLM configuration
    if same_llm:
        shared_agent_config: OpenAIGPTConfig = get_base_llm_config(
            await select_model("main LLM")
        )
        pro_agent_config = con_agent_config = shared_agent_config

        # Create feedback_agent_config by modifying shared_agent_config
        feedback_agent_config: OpenAIGPTConfig = OpenAIGPTConfig(
            chat_model=shared_agent_config.chat_model,
            min_output_tokens=shared_agent_config.min_output_tokens,
            max_output_tokens=shared_agent_config.max_output_tokens,
            temperature=0.2,  # Override temperature
            seed=shared_agent_config.seed,
        )
        metaphor_search_agent_config = feedback_agent_config
    else:
        pro_agent_config: OpenAIGPTConfig = get_base_llm_config(
            await select_model("for Pro Agent")
        )
        con_agent_config: OpenAIGPTConfig = get_base_llm_config(
            await select_model("for Con Agent")
        )
        feedback_agent_config: OpenAIGPTConfig = get_base_llm_config(
            await select_model("feedback"), temperature=0.2
        )
        metaphor_search_agent_config = feedback_agent_config

    system_messages: SystemMessages = load_system_messages(
        "examples/multi-agent-debate/system_messages.json"
    )
    LLM_DELEGATE_FLAG: bool = llm_delegate

    topic_name, pro_key, con_key, side = await select_topic_and_setup_side(
        LLM_DELEGATE_FLAG, system_messages
    )

    # Generate the system message
    metaphor_search_agent_system_message = (
        generate_metaphor_search_agent_system_message(system_messages, pro_key, con_key)
    )
    # pro_agent_system_message = "You are Pro. Start your response with 'Pro: "
    # + system_messages.messages[pro_key].message + DEFAULT_SYSTEM_MESSAGE_ADDITION
    pro_agent = create_custom_chat_agent(
        "Pro",
        pro_agent_config,
        system_messages.messages[pro_key].message + DEFAULT_SYSTEM_MESSAGE_ADDITION,
    )
    con_agent = create_custom_chat_agent(
        "Con",
        con_agent_config,
        system_messages.messages[con_key].message + DEFAULT_SYSTEM_MESSAGE_ADDITION,
    )
    feedback_agent = create_chat_agent(
        "Feedback", feedback_agent_config, FEEDBACK_AGENT_SYSTEM_MESSAGE
    )
    metaphor_search_agent = MetaphorSearchChatAgent(  # Use the subclass here
        ChatAgentConfig(
            llm=metaphor_search_agent_config,
            name="MetaphorSearch",
            system_message=metaphor_search_agent_system_message,
        )
    )

    logger.info("Pro, Con, feedback, and metaphor_search agents created.")

    # Determine user's side and assign user_agent and ai_agent based on user selection
    agents = {
        "pro": (pro_agent, con_agent, "Pro", "Con"),
        "con": (con_agent, pro_agent, "Con", "Pro"),
    }
    user_agent, ai_agent, user_side, ai_side = agents[side]
    logger.info(
        f"Starting debate on topic: {topic_name}, taking the {user_side} side. "
        f"LLM Delegate: {llm_delegate}"
    )

    logger.info(f"\n{user_side} Agent ({topic_name}):\n")

    # Determine if the debate is autonomous or the user input for one side
    if LLM_DELEGATE_FLAG:
        logger.info("Autonomous Debate Selected")
        interactive_setting = False
    else:
        logger.info("Manual Debate Selected with an AI Agent")
        interactive_setting = True
        user_input: str
        try:
            user_input_response = await cl.AskUserMessage(
                content="Your argument (or type 'f' for feedback, 'done' to end):",
                timeout=600,  # 10 minutes
            ).send()

            logger.info(f"Received user input response: {user_input_response}")

            if user_input_response and "output" in user_input_response:
                user_input = str(user_input_response["output"]).strip()
                logger.info(f"User input processed successfully: {user_input}")
                user_agent.llm = None  # User message without LLM completion
                user_agent.user_message = user_input
            else:
                logger.error("Response received but 'output' key is missing or empty.")
                raise TimeoutError(
                    "No valid response received for the user input question."
                )

        except TimeoutError as e:
            logger.error(str(e))
            # Handle timeout or invalid response gracefully

        # Assign the input to the user agent's attributes
        user_agent.llm = None  # User message without LLM completion
        user_agent.user_message = user_input

    # Set up langroid tasks and run the debate
    user_task = Task(user_agent, interactive=interactive_setting, restart=False)
    ai_task = Task(ai_agent, interactive=False, single_round=True)

    user_task.add_sub_task(ai_task)

    if not llm_delegate:
        ChainlitTaskCallbacks(user_task)
        await user_task.run_async(user_agent.user_message, turns=max_turns)
    else:
        CustomChainlitTaskCallbacks(user_task)
        await user_task.run_async("get started", turns=max_turns)

    # Determine the last agent based on turn count and alternation
    # Note: user_agent and ai_agent are dynamically set based on the chosen user_side
    last_agent = ai_agent if max_turns % 2 == 0 else user_agent

    await cl.Message(content="## Feedback and Debate Evaluation:").send()
    # Generate feedback summary and declare a winner using feedback agent

    if not last_agent.message_history:
        logger.warning("No agent message history found for the last agent")

    feedback_task = Task(
        feedback_agent,
        system_message=FEEDBACK_AGENT_SYSTEM_MESSAGE,
        interactive=False,
        single_round=True,
    )
    formatted_history = parse_and_format_message_history(last_agent.message_history)
    CustomChainlitTaskCallbacks(feedback_task)
    await feedback_task.run_async(
        formatted_history
    )  # Pass formatted history to the feedback agent

    metaphor_search: bool = await is_metaphor_search_key_set()

    if metaphor_search:
        metaphor_search_task = Task(
            metaphor_search_agent,
            system_message=metaphor_search_agent_system_message,
            interactive=False,
        )
        metaphor_search_agent.enable_message(MetaphorSearchTool)
        metaphor_search_agent.enable_message(DoneTool)
        CustomChainlitTaskCallbacks(metaphor_search_task)
        await metaphor_search_task.run_async("run the search")

        url_docs_ask_questions = await is_url_ask_question(topic_name)
        if url_docs_ask_questions:
            searched_urls = extract_urls(metaphor_search_agent.message_history)
            logger.info(searched_urls)
            ask_questions_agent = lr.agent.special.DocChatAgent(
                get_questions_agent_config(
                    searched_urls, feedback_agent_config.chat_model
                )
            )
            ask_questions_task = lr.Task(ask_questions_agent)
            CustomChainlitTaskCallbacks(ask_questions_task)
            await ask_questions_task.run_async()
</file>

<file path="examples/multi-agent-debate/main.py">
import logging
from typing import Any, List

import typer
from config import get_base_llm_config, get_global_settings, get_questions_agent_config
from models import SystemMessages, load_system_messages
from rich.prompt import Prompt
from system_messages import (
    DEFAULT_SYSTEM_MESSAGE_ADDITION,
    FEEDBACK_AGENT_SYSTEM_MESSAGE,
    generate_metaphor_search_agent_system_message,
)

# Import from utils.py
from utils import (
    extract_urls,
    is_llm_delegate,
    is_metaphor_search_key_set,
    is_same_llm_for_all_agents,
    is_url_ask_question,
    select_max_debate_turns,
    select_model,
    select_topic_and_setup_side,
)

import langroid as lr
from langroid import ChatDocument, Entity
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool
from langroid.agent.tools.orchestration import DoneTool
from langroid.language_models import OpenAIGPTConfig
from langroid.utils.logging import setup_logger


class MetaphorSearchChatAgent(ChatAgent):
    def handle_message_fallback(self, msg: str | ChatDocument) -> str | None:
        """Handle scenario where LLM did not generate any Tool"""
        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
            return f"""
            Have you presented pro and con arguments based on 
            your search results? If so, use the TOOL `{DoneTool.name()}` to indicate you're finished. 
            Otherwise, argue both sides and then send the `{DoneTool.name()}`
            """
        return None


# Initialize typer application
app = typer.Typer()

# set info logger
logger = setup_logger(__name__, level=logging.INFO, terminal=True)
logger.info("Starting multi-agent-debate")


def parse_and_format_message_history(message_history: List[Any]) -> str:
    """
    Parses and formats message history to exclude system messages
    and map roles to Pro/Con.

    Args:
        message_history (List[Any]): The full message history
        containing system, Pro, and Con messages.

    Returns:
        str: A formatted string with annotated Pro/Con messages.
    """
    annotated_history = []

    for msg in message_history:
        # Exclude system messages
        if msg.role == "system":
            continue

        # Map roles to Pro/Con
        if msg.role in ["pro", "user"]:  # User is treated as Pro in this context
            annotated_history.append(f"Pro: {msg.content}")
        elif msg.role in ["con", "assistant"]:  # Assistant is treated as Con
            annotated_history.append(f"Con: {msg.content}")

    return "\n".join(annotated_history)


def create_chat_agent(
    name: str, llm_config: OpenAIGPTConfig, system_message: str
) -> ChatAgent:
    """creates a ChatAgent with the given parameters.

    Args:
        name (str): The name of the agent.
        llm_config (OpenAIGPTConfig): The LLM configuration for the agent.
        system_message (str): The system message to guide the agent's LLM.

    Returns:
        ChatAgent: A configured ChatAgent instance.
    """
    return ChatAgent(
        ChatAgentConfig(
            llm=llm_config,
            name=name,
            system_message=system_message,
        )
    )


def run_debate() -> None:
    """Execute the main debate logic.

    Orchestrates the debate process, including setup, user input, LLM agent
    interactions, and final feedback. Handles both user-guided and LLM-
    delegated debates.

    This function:
    1. Loads global settings and the base LLM configurations.
    2. Prompts the user to confirm if they want to use same LLM for all agents.
    3. Prompts the user to select a debate topic and a side(Pro or Con).
    4. Sets up pro, con, and feedback agents.
    5. Runs the debate for a specified number of turns, either interactively
       or autonomously.
    6. Provides a feedback summary at the end.
    """

    global_settings = get_global_settings(nocache=True)
    lr.utils.configuration.set_global(global_settings)

    same_llm: bool = is_same_llm_for_all_agents()
    llm_delegate: bool = is_llm_delegate()
    max_turns: int = select_max_debate_turns()

    # Get base LLM configuration
    if same_llm:
        shared_agent_config: OpenAIGPTConfig = get_base_llm_config(
            select_model("main LLM")
        )
        pro_agent_config = con_agent_config = shared_agent_config

        # Create feedback_agent_config by modifying shared_agent_config
        feedback_agent_config: OpenAIGPTConfig = OpenAIGPTConfig(
            chat_model=shared_agent_config.chat_model,
            min_output_tokens=shared_agent_config.min_output_tokens,
            max_output_tokens=shared_agent_config.max_output_tokens,
            temperature=0.2,  # Override temperature
            seed=shared_agent_config.seed,
        )
        metaphor_search_agent_config = feedback_agent_config
    else:
        pro_agent_config: OpenAIGPTConfig = get_base_llm_config(
            select_model("for Pro Agent")
        )
        con_agent_config: OpenAIGPTConfig = get_base_llm_config(
            select_model("for Con Agent")
        )
        feedback_agent_config: OpenAIGPTConfig = get_base_llm_config(
            select_model("feedback"), temperature=0.2
        )
        metaphor_search_agent_config = feedback_agent_config

    system_messages: SystemMessages = load_system_messages(
        "examples/multi-agent-debate/system_messages.json"
    )
    topic_name, pro_key, con_key, side = select_topic_and_setup_side(system_messages)

    # Generate the system message
    metaphor_search_agent_system_message = (
        generate_metaphor_search_agent_system_message(system_messages, pro_key, con_key)
    )

    pro_agent = create_chat_agent(
        "Pro",
        pro_agent_config,
        system_messages.messages[pro_key].message + DEFAULT_SYSTEM_MESSAGE_ADDITION,
    )
    con_agent = create_chat_agent(
        "Con",
        con_agent_config,
        system_messages.messages[con_key].message + DEFAULT_SYSTEM_MESSAGE_ADDITION,
    )
    feedback_agent = create_chat_agent(
        "Feedback", feedback_agent_config, FEEDBACK_AGENT_SYSTEM_MESSAGE
    )
    metaphor_search_agent = MetaphorSearchChatAgent(  # Use the subclass here
        ChatAgentConfig(
            llm=metaphor_search_agent_config,
            name="MetaphorSearch",
            system_message=metaphor_search_agent_system_message,
        )
    )

    logger.info("Pro, Con, feedback, and metaphor_search agents created.")

    # Determine user's side and assign user_agent and ai_agent based on user selection
    agents = {
        "pro": (pro_agent, con_agent, "Pro", "Con"),
        "con": (con_agent, pro_agent, "Con", "Pro"),
    }
    user_agent, ai_agent, user_side, ai_side = agents[side]
    logger.info(
        f"Starting debate on topic: {topic_name}, taking the {user_side} side. "
        f"LLM Delegate: {llm_delegate}"
    )

    logger.info(f"\n{user_side} Agent ({topic_name}):\n")

    # Determine if the debate is autonomous or the user input for one side
    if llm_delegate:
        logger.info("Autonomous Debate Selected")
        interactive_setting = False
    else:
        logger.info("Manual Debate Selected with an AI Agent")
        interactive_setting = True
        user_input: str = Prompt.ask(
            "Your argument (or type 'f' for feedback, 'done' to end):"
        )
        user_agent.llm = None  # User message without LLM completion
        user_agent.user_message = user_input

    # Set up langroid tasks and run the debate
    user_task = Task(user_agent, interactive=interactive_setting, restart=False)
    ai_task = Task(ai_agent, interactive=False, single_round=True)
    user_task.add_sub_task(ai_task)
    if not llm_delegate:
        user_task.run(user_agent.user_message, turns=max_turns)
    else:
        user_task.run("get started", turns=max_turns)

    # Determine the last agent based on turn count and alternation
    # Note: user_agent and ai_agent are dynamically set based on the chosen user_side
    last_agent = ai_agent if max_turns % 2 == 0 else user_agent

    # Generate feedback summary and declare a winner using feedback agent
    if not last_agent.message_history:
        logger.warning("No agent message history found for the last agent")

    feedback_task = Task(feedback_agent, interactive=False, single_round=True)
    formatted_history = parse_and_format_message_history(last_agent.message_history)
    feedback_task.run(formatted_history)  # Pass formatted history to the feedback agent

    metaphor_search: bool = is_metaphor_search_key_set()

    if metaphor_search:
        metaphor_search_task = Task(metaphor_search_agent, interactive=False)
        metaphor_search_agent.enable_message(MetaphorSearchTool)
        metaphor_search_agent.enable_message(DoneTool)
        metaphor_search_task.run("run the search")

        url_docs_ask_questions = is_url_ask_question(topic_name)
        if url_docs_ask_questions:
            searched_urls = extract_urls(metaphor_search_agent.message_history)
            logger.info(searched_urls)
            ask_questions_agent = lr.agent.special.DocChatAgent(
                get_questions_agent_config(
                    searched_urls, feedback_agent_config.chat_model
                )
            )
            ask_questions_task = lr.Task(ask_questions_agent)
            ask_questions_task.run()


@app.command()
def main():
    """Main function and entry point for the Debate System"""
    run_debate()


if __name__ == "__main__":
    app()
</file>

<file path="examples/multi-agent-debate/models.py">
import json
import logging
from typing import Any, Dict

from pydantic import BaseModel
from langroid.utils.logging import setup_logger

logger = setup_logger(__name__, level=logging.INFO, terminal=True)


class Message(BaseModel):
    """Represents a single message with a topic and content.

    Attributes:
        topic (str): The topic of the message.
        message (str): The content of the message.
    """

    topic: str
    message: str


class SystemMessages(BaseModel):
    """Represents a collection of system messages.

    Attributes:
        messages (Dict[str, Message]): A dictionary where the key is the message
            identifier (e.g., 'pro_ai') and the value is a `Message` object.
    """

    messages: Dict[str, Message]


def load_system_messages(file_path: str) -> SystemMessages:
    """Load and validate system messages from a JSON file.

    Reads the JSON file containing system messages, maps each entry to a
    `Message` object, and wraps the result in a `SystemMessages` object.

    Args:
        file_path (str): The path to the JSON file containing system messages.

    Returns:
        SystemMessages: A `SystemMessages` object containing validated messages.

    Raises:
        IOError: If the file cannot be read or found.
        json.JSONDecodeError: If the JSON file is not properly formatted.
        Exception: For any other unexpected errors during processing.
    """
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            data: Any = json.load(f)
        # Map dictionaries to Message objects
        messages = {key: Message(**value) for key, value in data.items()}
        return SystemMessages(messages=messages)
    except FileNotFoundError as e:
        logger.error(f"File not found: {file_path}")
        raise IOError(f"Could not find the file: {file_path}") from e
    except json.JSONDecodeError as e:
        logger.error(f"Error decoding JSON file: {file_path}")
        raise json.JSONDecodeError(
            f"Invalid JSON format in file: {file_path}", e.doc, e.pos
        ) from e
    except Exception as e:
        logger.error(f"Unexpected error loading system messages: {e}")
        raise
</file>

<file path="examples/multi-agent-debate/README.md">
Debate System Using LLM Agents
==============================

Overview
--------
This project is a debate system powered by LLMs using Langroid, enabling structured debates on various topics 
such as AI in healthcare, education, intellectual property, and societal biases. 
The program creates and manages agents that represent opposing sides of a debate, 
interact with users, and provide constructive feedback based on established debate criteria.

New Topics and Pro and Con Side System messages can be manually configured by updating or modifying the 
system_messages.json File. 
"pro_ai": {
        "topic": "Your New TOPIC",
        "message": " YOUR Prompt"
    },
"con_ai": {
        "topic": "Your New TOPIC",
        "message": " YOUR CON or opposing Prompt"
        }

Features
--------
- Multiple Debate Topics:
  - AI in Healthcare
  - AI and Intellectual Property
  - AI and Societal Biases
  - AI as an Educator
- Agent-Based Interaction:
  - Pro and Con agents for each topic simulate structured debate arguments.
- Configurable to use different LLMs from OPENAI, Google, & Mistral: 
  -       1: gpt-4o
          2: gpt-4
          3: gpt-4o-mini
          4: gpt-4-turbo
          5: gpt-4-32k
          6: gpt-3.5-turbo-1106 
          7: Mistral: mistral:7b-instruct-v0.2-q8_0a
          8: Gemini:gemini-2.0-flash
          9: Gemini:gemini-1.5-flash
          10: Gemini:gemini-1.5-flash-8b
          11: Gemini:gemini-1.5-pro
- Feedback Mechanism:
  - Provides structured feedback on debate performance based on key criteria.
- Interactive or Autonomous Mode:
  - Users can either control interactions manually or let agents autonomously continue debates.

File Structure
--------------
- main.py: The entry point of the application. Initializes the system, configures agents, and starts the debate loop.
- config.py: Provides functions for configuring global settings and LLM-specific parameters.
- model.py: Pydantic model for system_messages.json
- system_messages.json: Topic Titles and system_messages for pro and con agents. You can add more topics and their
respective pro and con system messages here. The system_messages has a statement: 
"Limit responses to MAXIMUM 2 points expressed as single sentences." Please change or delete it for a realistic debate. 
- system_message.py: Global system messages
- utils.py: User Prompts and other helper functions
- generation_config_models.py: pydantic model for generation_config.json
- generation_config.json: LLM generation parameters
- main_chainlit.py: The entry point of the Chainlit version of the application
- chainlit_utils.py: Relevant chainlit utility functions.
The system dynamically updates user selection with the topics from this file. 

Getting Started
---------------
Prerequisites
1. Python 3.8+
2. Langroid Framework: Install Langroid with necessary dependencies:
   pip install "langroid[litellm]"
3. Setup the following env variables in the .env File in the root of your repo
or set them on your terminal.
       export OPENAI_API_KEY=OPEN AI KEY
       export GEMINI_API_KEY=GEMiNi API KEY
       export METAPHOR_API_KEY=METAPHOR_API_KEY
4. Please read the following page for more information:
   https://langroid.github.io/langroid/quick-start/setup/

Usage
-----
Run the CLI Application
Start the application from the root of the langroid repo with:
   python examples/multi-agent-debate/main.py

Options
- Debug Mode: Run the program with debug logs for detailed output.
  python examples/multi-agent-debate/main.py --debug
- Disable Caching: Avoid using cached responses for LLM interactions.
  python examples/multi-agent-debate/main.py --nocache

Run the Chainlit App
chainlit run examples/multi-agent-debate/main_chainlit.py


Interaction
1. Decide if you want to you use same LLM for all agents or different ones
2. Decide if you want autonomous debate between AI Agents or user vs. AI Agent. 
3. Select a debate topic.
4. Choose your side (Pro or Con).
5. Engage in a debate by providing arguments and receiving responses from agents.
6. Request feedback at any time by typing `f`.
7. Decide if you want the Metaphor Search to run to find Topic relevant web links
   and summarize them. 
8. Decide if you want to chat with the documents extracted from URLs found to learn more about the Topic.
9. End the debate manually by typing "done". If you decide to chat with the documents, you can end session
by typing `x`

Feedback Criteria
-----------------
The feedback mechanism evaluates debates based on:
1. Clash of Values
2. Argumentation
3. Cross-Examination
4. Rebuttals
5. Persuasion
6. Technical Execution
7. Adherence to Debate Etiquette
8. Final Focus

License
-------
This project is licensed under the MIT License.
</file>

<file path="examples/multi-agent-debate/system_messages.py">
from langroid.agent.tools.metaphor_search_tool import MetaphorSearchTool
from langroid.agent.tools.orchestration import DoneTool

DEFAULT_SYSTEM_MESSAGE_ADDITION = """
            DO NOT REPEAT ARGUMENTS THAT HAVE BEEN PREVIOUSLY GENERATED 
            AND CAN BE SEEN IN THE DEBATE HISTORY PROVIDED. 
            """
FEEDBACK_AGENT_SYSTEM_MESSAGE = """  
            You are an expert and experienced judge specializing in Lincoln-Douglas style debates. 
            Your goal is to evaluate the debate thoroughly based on the following criteria:
            1. Clash of Values: Assess how well each side upholds their stated value (e.g., justice, morality) 
               and how effectively they compare and prioritize values.
            2. Argumentation: Evaluate the clarity, organization, and logical soundness of each side's case structure, 
               contentions, and supporting evidence.
            3. Cross-Examination: Judge the effectiveness of questioning and answering during cross-examination.
            4. Rebuttals: Analyze how well each side refutes their opponent's arguments.
            5. Persuasion: Assess communication quality, tone, rhetorical effectiveness, and emotional/ethical appeals.
            6. Technical Execution: Identify if major arguments were addressed or dropped and check consistency.
            7. Debate Etiquette: Evaluate professionalism, respect, and demeanor.
            8. Final Focus: Judge the strength of closing speeches, how well they summarize the case, 
            and justify a winner.
            Provide constructive feedback for each debater, 
            summarizing their performance and declaring a winner with justification.   
            """
METAPHOR_SEARCH_AGENT_SYSTEM_MESSAGE_TEMPLATE = """
            There are 2 STEPs. Your Goal is to execute both of them. 
            STEP 1:  Run MetaphorSearchTool

            Use the TOOL {metaphor_tool_name} to search the web for 5 references for Pro: {pro_message}
            and Con: {con_message}.     
            YOUR GOAL IS TO FIND GOOD REFERENCES FOR BOTH SIDES OF A DEBATE. 
            Be very CONCISE in your responses, use 5-7 sentences. 
            show me the SOURCE(s) and EXTRACT(s) and summary
            in this format:

            <your answer here>
            Here are additional references using Metaphor Search to improve your knowledge of the subject:

            M1: SOURCE: https://journalofethics.ama-assn.org/article/should-artificial-intelligence-augment-
            medical-decision-making-case-autonomy-algorithm/2018-09
            EXTRACT: Discusses the ethical implications of AI in medical decision-making and 
            the concept of an autonomy algorithm.
            SUMMARY: This article explores the ethical considerations of integrating AI into medical decision-making 
            processes, emphasizing the need for autonomy and ethical oversight.

            M2: SOURCE: ...
            EXTRACT: ...
            SUMMARY:

            DO NOT MAKE UP YOUR OWN SOURCES; ONLY USE SOURCES YOU FIND FROM A WEB SEARCH. 
            ENSURE STEP 1 IS COMPLETED BEFORE STARTING STEP 2

            STEP 2: Argue Pro and Con Cases
            As an expert debater, your goal is to eloquently argue for both the Pro and Con cases
            using the references from web-search SOURCES generated in Step 1 and properly cite the Sources in BRACKETS 
            (e.g., [SOURCE])
            Write at least 5 sentences for each side.

            ENSURE BOTH STEP 1 and 2 are completed. 
            After all STEPs are completed, use the `{done_tool_name}` tool to end the session   
            """


def generate_metaphor_search_agent_system_message(system_messages, pro_key, con_key):
    return METAPHOR_SEARCH_AGENT_SYSTEM_MESSAGE_TEMPLATE.format(
        metaphor_tool_name=MetaphorSearchTool.name(),
        pro_message=system_messages.messages[pro_key].message,
        con_message=system_messages.messages[con_key].message,
        done_tool_name=DoneTool.name(),
    )
</file>

<file path="examples/multi-agent-debate/utils.py">
import logging
import re
from typing import List, Literal, Optional, Tuple

from models import SystemMessages
from rich.prompt import Confirm, Prompt

from langroid.utils.logging import setup_logger

DEFAULT_TURN_COUNT = 2

# set info logger
logger = setup_logger(__name__, level=logging.INFO, terminal=True)


def extract_topics(system_messages: SystemMessages) -> List[Tuple[str, str, str]]:
    """Extract unique debate topics from the SystemMessages object.

    Processes the `SystemMessages` object to identify debate topics by pairing
    `pro_` and `con_` keys. Ensures each topic is represented only once.

    Args:
        system_messages (SystemMessages): The object containing system messages
            with `pro_` and `con_` topic keys.

    Returns:
        List[Tuple[str, str, str]]: A list of tuples,
        where each tuple contains:
            - topic_name (str): The name of the debate topic.
            - pro_key (str): The key for the pro side of the debate.
            - con_key (str): The key for the con side of the debate.
    """
    topics: List[Tuple[str, str, str]] = []
    for key, message in system_messages.messages.items():
        # Process only "pro_" keys to avoid duplicates
        if key.startswith("pro_"):
            con_key = key.replace("pro_", "con_", 1)  # Match "con_" dynamically
            if con_key in system_messages.messages:  # Ensure "con_" exists
                topics.append((message.topic, key, con_key))
    return topics


def select_model(config_agent_name: str) -> str:
    """
    Prompt the user to select an OpenAI or Gemini model
    for the specified agent.

    This function prompts the user to select an option from
    a list of available models.
    The user's input corresponds to a predefined choice, which is
    then returned as a string representing the selected option.

    Args:
        config_agent_name (str): The name of the agent being configured,
        used in the prompt to personalize the message.

    Returns:
        str: The user's selected option as a string, corresponding to one of the
             predefined model choices (e.g., "1", "2", ..., "10").
    """
    return Prompt.ask(
        f"Select a Model for {config_agent_name}:\n"
        "1: gpt-4o\n"
        "2: gpt-4\n"
        "3: gpt-4o-mini\n"
        "4: gpt-4-turbo\n"
        "5: gpt-4-32k\n"
        "6: gpt-3.5-turbo-1106\n"
        "7: Mistral: mistral:7b-instruct-v0.2-q8_0a\n"
        "8: Gemini: gemini-2.0-flash\n"
        "8: Gemini: gemini-1.5-flash\n"
        "9: Gemini: gemini-1.5-flash-8b\n"
        "10: Gemini: gemini-1.5-pro\n",
        choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"],
        default="1",
    )


def select_debate_topic(system_messages: SystemMessages) -> Optional[tuple]:
    """Prompt the user to select a debate topic from SystemMessages.

    Dynamically loads debate topics from the SystemMessages object, displays
    the options to the user, and prompts them to select a topic.

    Args:
        system_messages (SystemMessages): The object containing debate topics.

    Returns:
        Optional[tuple]: A tuple containing:
            - topic_name (str): The selected topic's name.
            - pro_key (str): The key for the pro side of the debate.
            - con_key (str): The key for the con side of the debate.
            Returns None if no topics are available or an error occurs.
    """
    topics = extract_topics(system_messages)

    if not topics:
        logger.error("No topics found in the JSON file.")
        return None

    # Prepare topic choices for user selection
    topic_choices = "\n".join(
        [f"{i + 1}. {topic[0]}" for i, topic in enumerate(topics)]
    )
    user_input = Prompt.ask(
        f"Select a debate topic:\n{topic_choices}",
        choices=[str(i + 1) for i in range(len(topics))],
        default="1",
    )
    topic_index = int(user_input) - 1

    selected_topic = topics[topic_index]
    logger.info(f"Selected topic: {selected_topic[0]}")
    return selected_topic


def select_side(topic_name: str) -> Literal["pro", "con"]:
    """Prompt the user to select their side in the debate.

    Presents the user with a choice to debate on either the pro or con side
    of the given topic.

    Args:
        topic_name (str): The name of the debate topic.

    Returns:
        Literal["pro", "con"]: The selected debate side.
    """
    side = Prompt.ask(
        f"Which side would you like to debate on?\n1. Pro-{topic_name}\n2. "
        f"Con-{topic_name}",
        choices=["1", "2"],
        default="1",
    )
    return "pro" if side == "1" else "con"


def select_topic_and_setup_side(
    system_messages: SystemMessages,
) -> Tuple[str, str, str, str]:
    """Prompt the user to select a debate topic and sets up the respective side.

    This function handles the user interaction for selecting a debate topic and the side
    (Pro or Con) they want to argue. It validates that a topic is selected and raises an
    exception if the topic is not available.

    Args:
        system_messages (SystemMessages): The object containing system messages with respective
                                          debate topics.

    Returns:
        Tuple[str, str, str, str]: A tuple containing:
            - topic_name (str): The name of the selected debate topic.
            - pro_key (str): The key for the Pro side of the selected topic.
            - con_key (str): The key for the Con side of the selected topic.
            - side (str): The user's selected side, either "pro" or "con".

    Raises:
        ValueError: If no topic is selected or no topics are available in the provided
                    `system_messages`.
    """
    selected_topic_tuple = select_debate_topic(system_messages)
    if not selected_topic_tuple:
        logger.error("No topic selected. Exiting.")
        raise ValueError("No topic selected.")

    topic_name, pro_key, con_key = selected_topic_tuple
    side = select_side(topic_name)
    return topic_name, pro_key, con_key, side


def is_llm_delegate() -> bool:
    """Prompt the user to decide on LLM delegation.

    Asks the user whether the LLM should autonomously continue the debate
    without requiring user input.

    Returns:
        bool: True if the user chooses LLM delegation, otherwise return False.
    """
    return Confirm.ask(
        "Should the Pro and Con agents debate autonomously?",
        default=False,
    )


def is_metaphor_search_key_set() -> bool:
    """Prompt the user confirmation about metaphorSearch API keys.

    Asks the user to confirm the metaphorSearch API keys.

    Returns:
        bool: True if the user chooses LLM delegation, otherwise return False.
    """
    return Confirm.ask(
        "Do you have an API Key for Metaphor Search?",
        default=False,
    )


def is_same_llm_for_all_agents() -> bool:
    """Prompt the user to decide if same LLM should be used for all agents.

    Asks the user whether the same LLM should be configured for all agents.

    Returns:
        bool: True if the user chooses same LLM for all agents, otherwise return False.
    """
    # Ask the user if they want to use the same LLM configuration for all agents
    return Confirm.ask(
        "Do you want to use the same LLM for all agents?",
        default=True,
    )


def select_max_debate_turns() -> int:
    # Prompt for number of debate turns
    while True:
        max_turns = Prompt.ask(
            "How many turns should the debate continue for?",
            default=str(DEFAULT_TURN_COUNT),
        )
        try:
            return int(max_turns)
        except ValueError:
            return DEFAULT_TURN_COUNT


def extract_urls(message_history):
    """
    Extracts all URLs from the given message history content and returns them in the format [url1, url2, ..., urln].

    Parameters:
        message_history (list): A list of LLMMessage objects containing message history.

    Returns:
        str: A string representation of a list of URLs.
    """
    # Extract content only from non-system messages
    content = " ".join(
        message.content
        for message in message_history
        if hasattr(message, "content") and message.content and message.role != "system"
    )

    # Extract URLs from content
    urls = re.findall(r"https?://\S+", content)
    return urls  # Return the list of URLs directly


def is_url_ask_question(topic_name: str) -> bool:
    """Prompt the user to decide to ask questions from the searched URL docs.

    Asks the user whether they want to ask questions from the searched URL docs?

    Returns:
        bool: True if the user chooses to ask questions from searched url docs., otherwise return False.
    """
    return Confirm.ask(
        f"Would you like to Chat with documents found through Search for more information on the {topic_name}",
        default=False,
    )
</file>

<file path="examples/portkey/portkey_advanced_features.py">
#!/usr/bin/env python3

"""
Advanced Portkey example showing observability, caching, retries, and metadata.

This example demonstrates:
- Advanced Portkey configuration with all features
- Request tracing and metadata
- Caching and retry strategies
- Custom headers for observability

Run with: python portkey_advanced_features.py
"""

import os
import uuid
from typing import Optional

import langroid as lr
import langroid.language_models as lm
from langroid.language_models.provider_params import PortkeyParams


def check_env_var(var_name: str) -> Optional[str]:
    """Check if environment variable is set and return its value."""
    value = os.getenv(var_name)
    if not value:
        print(f"⚠️  Warning: {var_name} not set in environment")
        return None
    return value


def create_advanced_portkey_llm(portkey_api_key: str, user_id: str) -> lm.OpenAIGPT:
    """Create an advanced Portkey-enabled LLM with all features."""

    # Generate unique trace ID for this session
    trace_id = f"trace-{uuid.uuid4().hex[:8]}"

    config = lm.OpenAIGPTConfig(
        chat_model="portkey/openai/gpt-4o-mini",
        portkey_params=PortkeyParams(
            api_key=portkey_api_key,
            # Observability features
            trace_id=trace_id,
            metadata={
                "user_id": user_id,
                "app": "langroid-advanced-example",
                "version": "1.0",
                "environment": "demo",
            },
            # Retry configuration
            retry={"max_retries": 3, "backoff": "exponential", "jitter": True},
            # Caching configuration
            cache={
                "enabled": True,
                "ttl": 3600,  # 1 hour
                "namespace": "langroid-demo",
            },
            cache_force_refresh=False,
            # User tracking
            user=user_id,
            organization="langroid-demo-org",
            # Custom headers for additional tracking
            custom_headers={
                "x-session-id": f"session-{uuid.uuid4().hex[:8]}",
                "x-demo-type": "advanced-features",
                "x-langroid-version": (
                    lr.__version__ if hasattr(lr, "__version__") else "unknown"
                ),
            },
        ),
        max_output_tokens=200,
        temperature=0.3,
    )

    return lm.OpenAIGPT(config)


def demonstrate_caching(llm: lm.OpenAIGPT):
    """Demonstrate Portkey's caching capabilities."""
    print("\n🧠 Testing Caching Capabilities")
    print("=" * 50)

    question = "What are the three laws of robotics by Isaac Asimov?"

    print("🔄 First request (should hit the API)...")
    response1 = llm.chat(question)
    print(f"✅ Response: {response1.message[:100]}...")
    print(f"📊 Cached: {response1.cached}")
    if response1.usage:
        print(f"📊 Tokens: {response1.usage.total_tokens}")

    print("\n🔄 Second identical request (should hit cache)...")
    response2 = llm.chat(question)
    print(f"✅ Response: {response2.message[:100]}...")
    print(f"📊 Cached: {response2.cached}")
    if response2.usage:
        print(f"📊 Tokens: {response2.usage.total_tokens}")


def demonstrate_metadata_tracking(llm: lm.OpenAIGPT, user_id: str):
    """Demonstrate request tracking with metadata."""
    print("\n📊 Testing Metadata and Tracking")
    print("=" * 50)

    questions = [
        "What is machine learning?",
        "Explain neural networks briefly.",
        "What is the difference between AI and ML?",
    ]

    for i, question in enumerate(questions, 1):
        print(f"\n🔍 Question {i}: {question}")

        # Create a new LLM instance with updated metadata for each question
        trace_id = f"trace-q{i}-{uuid.uuid4().hex[:6]}"

        config = lm.OpenAIGPTConfig(
            chat_model="portkey/openai/gpt-4o-mini",
            portkey_params=PortkeyParams(
                api_key=llm.config.portkey_params.api_key,
                trace_id=trace_id,
                metadata={
                    "user_id": user_id,
                    "question_number": i,
                    "question_category": "AI/ML basics",
                    "timestamp": str(uuid.uuid4()),  # Unique per request
                },
                user=user_id,
                custom_headers={
                    "x-question-id": f"q-{i}",
                    "x-session-type": "educational",
                },
            ),
            max_output_tokens=150,
            temperature=0.3,
        )

        question_llm = lm.OpenAIGPT(config)
        response = question_llm.chat(question)

        print(f"✅ Response: {response.message[:80]}...")
        print(f"🏷️  Trace ID: {trace_id}")


def demonstrate_error_handling():
    """Demonstrate error handling and retry behavior."""
    print("\n⚠️  Testing Error Handling")
    print("=" * 50)

    try:
        # Create config with invalid model to test error handling
        config = lm.OpenAIGPTConfig(
            chat_model="portkey/openai/invalid-model-name",
            portkey_params=PortkeyParams(
                api_key=os.getenv("PORTKEY_API_KEY", ""),
                retry={"max_retries": 2, "backoff": "linear"},
                metadata={"test_type": "error_handling"},
            ),
        )

        error_llm = lm.OpenAIGPT(config)
        response = error_llm.chat("This should fail")
        print(f"Unexpected success: {response.message}")

    except Exception as e:
        print(f"✅ Expected error caught: {type(e).__name__}")
        print(f"   Error details: {str(e)[:100]}...")


def main():
    """Main function demonstrating advanced Portkey features."""
    print("🚀 Portkey Advanced Features Example")
    print("=" * 45)

    # Check for required environment variables
    portkey_api_key = check_env_var("PORTKEY_API_KEY")
    if not portkey_api_key:
        print("❌ PORTKEY_API_KEY is required. Please set it in your environment.")
        return

    openai_api_key = check_env_var("OPENAI_API_KEY")
    if not openai_api_key:
        print("❌ OPENAI_API_KEY is required for this example.")
        return

    print("✅ All required API keys found")

    # Generate a unique user ID for this session
    user_id = f"user-{uuid.uuid4().hex[:8]}"
    print(f"🆔 Demo User ID: {user_id}")

    # Create advanced LLM configuration
    try:
        llm = create_advanced_portkey_llm(portkey_api_key, user_id)
        print("✅ Advanced Portkey LLM created successfully")

        # Demonstrate different features
        demonstrate_caching(llm)
        demonstrate_metadata_tracking(llm, user_id)
        demonstrate_error_handling()

    except Exception as e:
        print(f"❌ Failed to create advanced LLM: {str(e)}")
        return

    print("\n🎉 Advanced Portkey features example completed!")
    print("\n💡 Next steps:")
    print(
        "   - View detailed request logs in Portkey dashboard: https://app.portkey.ai"
    )
    print(
        "   - Filter by trace IDs, user IDs, or metadata to analyze specific requests"
    )
    print("   - Try the multi-provider example: portkey_multi_provider.py")
    print(f"\n🔍 Your demo user ID: {user_id}")
    print("   Use this to filter requests in the Portkey dashboard")


if __name__ == "__main__":
    main()
</file>

<file path="examples/portkey/portkey_basic_chat.py">
#!/usr/bin/env python3

"""
Basic Portkey example showing how to use different AI providers through Portkey's gateway.

This example demonstrates:
- Basic Portkey configuration
- Switching between different AI providers
- Automatic API key resolution

Run with: python portkey_basic_chat.py
"""

import os
from typing import Optional

import langroid.language_models as lm
from langroid.language_models.provider_params import PortkeyParams


def check_env_var(var_name: str) -> Optional[str]:
    """Check if environment variable is set and return its value."""
    value = os.getenv(var_name)
    if not value:
        print(f"⚠️  Warning: {var_name} not set in environment")
        return None
    return value


def create_portkey_llm(provider: str, model: str, portkey_api_key: str) -> lm.OpenAIGPT:
    """Create a Portkey-enabled LLM configuration."""
    config = lm.OpenAIGPTConfig(
        chat_model=f"portkey/{provider}/{model}",
        portkey_params=PortkeyParams(
            api_key=portkey_api_key,
        ),
        max_output_tokens=150,
        temperature=0.7,
    )
    return lm.OpenAIGPT(config)


def test_provider(llm: lm.OpenAIGPT, provider_name: str):
    """Test a specific provider with a simple question."""
    print(f"\n🔮 Testing {provider_name}...")
    print("=" * 50)

    question = "What is the capital of France? Answer in one sentence."

    try:
        response = llm.chat(question)
        print(f"✅ {provider_name} Response:")
        print(f"   {response.message}")

        if response.usage:
            print(f"   📊 Tokens: {response.usage.total_tokens}")

    except Exception as e:
        print(f"❌ {provider_name} Error: {str(e)}")


def main():
    """Main function demonstrating Portkey basic usage."""
    print("🚀 Portkey Basic Chat Example")
    print("=" * 40)

    # Check for required environment variables
    portkey_api_key = check_env_var("PORTKEY_API_KEY")
    if not portkey_api_key:
        print("❌ PORTKEY_API_KEY is required. Please set it in your environment.")
        return

    print("✅ Portkey API key found")

    # Test different providers through Portkey
    providers_to_test = []

    # Check which provider keys are available
    if os.getenv("OPENAI_API_KEY"):
        providers_to_test.append(("OpenAI", "openai", "gpt-4o-mini"))
        print("✅ OpenAI API key found")

    if os.getenv("ANTHROPIC_API_KEY"):
        providers_to_test.append(("Anthropic", "anthropic", "claude-3-haiku-20240307"))
        print("✅ Anthropic API key found")

    if os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY"):
        providers_to_test.append(("Google Gemini", "google", "gemini-2.0-flash-lite"))
        print("✅ Google/Gemini API key found")

    if not providers_to_test:
        print("❌ No provider API keys found. Please set at least one of:")
        print("   - OPENAI_API_KEY")
        print("   - ANTHROPIC_API_KEY")
        print("   - GOOGLE_API_KEY or GEMINI_API_KEY")
        return

    print(f"\n🎯 Testing {len(providers_to_test)} provider(s) through Portkey...")

    # Test each available provider
    for provider_display_name, provider, model in providers_to_test:
        try:
            llm = create_portkey_llm(provider, model, portkey_api_key)
            test_provider(llm, provider_display_name)
        except Exception as e:
            print(f"❌ Failed to create {provider_display_name} LLM: {str(e)}")

    print("\n🎉 Portkey basic chat example completed!")
    print("\n💡 Next steps:")
    print("   - Try the advanced features example: portkey_advanced_features.py")
    print("   - View your requests in the Portkey dashboard: https://app.portkey.ai")


if __name__ == "__main__":
    main()
</file>

<file path="examples/portkey/portkey_multi_provider.py">
#!/usr/bin/env python3

"""
Multi-provider Portkey example showing how to compare responses across different AI providers.

This example demonstrates:
- Using multiple providers through Portkey
- Comparing response quality and characteristics
- Provider-specific configurations
- Fallback strategies

Run with: python portkey_multi_provider.py
"""

import os
import time
from typing import List, Optional, Tuple

import langroid.language_models as lm
from langroid.language_models.provider_params import PortkeyParams


def check_env_var(var_name: str) -> Optional[str]:
    """Check if environment variable is set and return its value."""
    value = os.getenv(var_name)
    if not value:
        print(f"⚠️  Warning: {var_name} not set in environment")
        return None
    return value


def create_provider_llm(
    provider: str, model: str, portkey_api_key: str, temperature: float = 0.7
) -> Tuple[lm.OpenAIGPT, str]:
    """Create a Portkey-enabled LLM for a specific provider."""
    config = lm.OpenAIGPTConfig(
        chat_model=f"portkey/{provider}/{model}",
        portkey_params=PortkeyParams(
            api_key=portkey_api_key,
            metadata={
                "provider": provider,
                "model": model,
                "demo": "multi-provider-comparison",
            },
            user="multi-provider-demo",
        ),
        max_output_tokens=200,
        temperature=temperature,
    )

    display_name = f"{provider.title()} ({model})"
    return lm.OpenAIGPT(config), display_name


def test_providers_on_question(
    providers: List[Tuple[lm.OpenAIGPT, str]], question: str
):
    """Test all providers on the same question and compare responses."""
    print(f"\n❓ Question: {question}")
    print("=" * 80)

    responses = []

    for llm, display_name in providers:
        print(f"\n🤖 {display_name}:")
        print("-" * 40)

        try:
            start_time = time.time()
            response = llm.chat(question)
            end_time = time.time()

            print(f"📝 Response: {response.message}")
            print(f"⏱️  Time: {end_time - start_time:.2f}s")

            if response.usage:
                print(f"📊 Tokens: {response.usage.total_tokens}")

            responses.append((display_name, response.message, response.usage))

        except Exception as e:
            print(f"❌ Error: {str(e)}")
            responses.append((display_name, f"Error: {str(e)}", None))

    return responses


def demonstrate_creative_tasks(providers: List[Tuple[lm.OpenAIGPT, str]]):
    """Test providers on creative tasks to see different capabilities."""
    print("\n🎨 Creative Tasks Comparison")
    print("=" * 50)

    creative_questions = [
        "Write a haiku about artificial intelligence.",
        "Explain quantum computing using a food analogy.",
        "Create a short story opening with exactly 50 words.",
    ]

    for question in creative_questions:
        test_providers_on_question(providers, question)


def demonstrate_analytical_tasks(providers: List[Tuple[lm.OpenAIGPT, str]]):
    """Test providers on analytical tasks."""
    print("\n🧮 Analytical Tasks Comparison")
    print("=" * 50)

    analytical_questions = [
        "What are the pros and cons of renewable energy?",
        "Explain the causes of inflation in simple terms.",
        "Compare machine learning and traditional programming.",
    ]

    for question in analytical_questions:
        test_providers_on_question(providers, question)


def demonstrate_fallback_strategy(portkey_api_key: str):
    """Demonstrate a simple fallback strategy across providers."""
    print("\n🔄 Fallback Strategy Demo")
    print("=" * 50)

    # Define providers in order of preference
    fallback_providers = [
        ("openai", "gpt-4o-mini", "OPENAI_API_KEY"),
        ("anthropic", "claude-3-haiku-20240307", "ANTHROPIC_API_KEY"),
        ("google", "gemini-2.0-flash-lite", "GOOGLE_API_KEY"),
    ]

    question = "What is the meaning of life in one sentence?"

    for provider, model, env_var in fallback_providers:
        if os.getenv(env_var):
            print(f"\n🎯 Trying {provider.title()}...")
            try:
                llm, display_name = create_provider_llm(
                    provider, model, portkey_api_key, temperature=0.5
                )
                response = llm.chat(question)
                print(f"✅ Success with {display_name}")
                print(f"📝 Response: {response.message}")
                return  # Success, stop trying

            except Exception as e:
                print(f"❌ {display_name} failed: {str(e)}")
                print("🔄 Trying next provider...")
        else:
            print(f"⏭️  Skipping {provider.title()} (API key not available)")

    print("❌ All providers failed!")


def main():
    """Main function demonstrating multi-provider usage."""
    print("🚀 Portkey Multi-Provider Example")
    print("=" * 45)

    # Check for required environment variables
    portkey_api_key = check_env_var("PORTKEY_API_KEY")
    if not portkey_api_key:
        print("❌ PORTKEY_API_KEY is required. Please set it in your environment.")
        return

    print("✅ Portkey API key found")

    # Collect available providers
    providers = []

    if os.getenv("OPENAI_API_KEY"):
        try:
            llm, name = create_provider_llm("openai", "gpt-4o-mini", portkey_api_key)
            providers.append((llm, name))
            print("✅ OpenAI provider ready")
        except Exception as e:
            print(f"⚠️  OpenAI setup failed: {e}")

    if os.getenv("ANTHROPIC_API_KEY"):
        try:
            llm, name = create_provider_llm(
                "anthropic", "claude-3-haiku-20240307", portkey_api_key
            )
            providers.append((llm, name))
            print("✅ Anthropic provider ready")
        except Exception as e:
            print(f"⚠️  Anthropic setup failed: {e}")

    if os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY"):
        try:
            llm, name = create_provider_llm(
                "google", "gemini-2.0-flash-lite", portkey_api_key
            )
            providers.append((llm, name))
            print("✅ Google/Gemini provider ready")
        except Exception as e:
            print(f"⚠️  Google/Gemini setup failed: {e}")

    if len(providers) < 2:
        print("\n⚠️  This example works best with at least 2 providers.")
        print("   Please set API keys for multiple providers:")
        print("   - OPENAI_API_KEY")
        print("   - ANTHROPIC_API_KEY")
        print("   - GOOGLE_API_KEY or GEMINI_API_KEY")

        if len(providers) == 0:
            print("❌ No providers available. Exiting.")
            return

    print(f"\n🎯 Ready to compare {len(providers)} provider(s)")

    # Run comparisons
    if len(providers) >= 2:
        demonstrate_creative_tasks(providers)
        demonstrate_analytical_tasks(providers)

    # Always demonstrate fallback (works with 1+ providers)
    demonstrate_fallback_strategy(portkey_api_key)

    print("\n🎉 Multi-provider example completed!")
    print("\n💡 Analysis tips:")
    print("   - Different providers may excel at different types of tasks")
    print("   - Response styles and lengths can vary significantly")
    print("   - Use Portkey dashboard to analyze performance metrics")
    print("   - Consider cost, speed, and quality when choosing providers")
    print("\n🔍 View detailed comparisons at: https://app.portkey.ai")


if __name__ == "__main__":
    main()
</file>

<file path="examples/portkey/README.md">
# Portkey Examples

This folder contains examples demonstrating how to use [Portkey](https://portkey.ai) with Langroid for enhanced LLM gateway functionality and observability.

## Prerequisites

Before running any examples, make sure you've installed Langroid as usual.

At minimum, have these environment variables set up in your `.env` file or environment:
```bash
PORTKEY_API_KEY=your_portkey_api_key_here
OPENAI_API_KEY=your_openai_api_key_here  # or any provider's key
ANTHROPIC_API_KEY=your_anthropic_key_here  # if using Anthropic
```

### 1. Portkey Basic Chat (`portkey_basic_chat.py`)

Demonstrates basic chat functionality with Portkey:
- Uses Portkey as a gateway to different AI providers
- Shows automatic provider API key resolution
- Demonstrates model switching across providers

```python
# Run the example from root of repo, after activating your virtual environment with uv:
uv run examples/portkey/portkey_basic_chat.py
```

### 2. Portkey Advanced Features (`portkey_advanced_features.py`)

Shows how to use Portkey's advanced features:
- Virtual keys for provider management
- Caching and retry configurations
- Request tracing and metadata
- Custom headers for observability

```python
# Run the example
uv run examples/portkey/portkey_advanced_features.py
```

### 3. Portkey Multi-Provider Example (`portkey_multi_provider.py`)

Showcases Portkey's ability to switch between providers:
- Compares responses from different providers
- Demonstrates fallback strategies
- Shows how to use virtual keys for different models

```python
# Run the example
uv run examples/portkey/portkey_multi_provider.py
```

## Using Portkey

### Basic Configuration

Portkey can route requests to any AI provider through a unified API:

```python
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.language_models.provider_params import PortkeyParams

# Configure for OpenAI via Portkey
config = OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o-mini",
    portkey_params=PortkeyParams(
        api_key="your-portkey-api-key",  # Or use PORTKEY_API_KEY env var
    )
)

# Configure for Anthropic via Portkey
config = OpenAIGPTConfig(
    chat_model="portkey/anthropic/claude-3-sonnet-20240229",
    portkey_params=PortkeyParams(
        api_key="your-portkey-api-key",
    )
)
```

### Advanced Features

Portkey provides powerful gateway features:

```python
from langroid.language_models.provider_params import PortkeyParams

# Advanced configuration with observability
params = PortkeyParams(
    api_key="your-portkey-api-key",
    virtual_key="vk-your-virtual-key",  # For provider abstraction
    trace_id="trace-123",               # For request tracing
    metadata={"user": "john", "app": "langroid"},  # Custom metadata
    retry={"max_retries": 3, "backoff": "exponential"},  # Retry config
    cache={"enabled": True, "ttl": 3600},  # Caching config
    cache_force_refresh=False,          # Cache control
    user="user-123",                    # User identifier
    organization="org-456",             # Organization identifier
    custom_headers={                    # Additional custom headers
        "x-custom-header": "value"
    }
)

config = OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o",
    portkey_params=params
)
```

### Supported Providers

Portkey supports many AI providers:

```python
# OpenAI
chat_model="portkey/openai/gpt-4o"

# Anthropic
chat_model="portkey/anthropic/claude-3-5-sonnet-20241022"

# Google Gemini
chat_model="portkey/google/gemini-2.0-flash-lite"

# Cohere
chat_model="portkey/cohere/command-r-plus"

# Many more providers available through Portkey
```

### Environment Variables

Portkey integration automatically resolves API keys from environment variables:

```bash
# Portkey API key
PORTKEY_API_KEY=your_portkey_api_key

# Provider API keys (used for actual model calls)
OPENAI_API_KEY=your_openai_key
ANTHROPIC_API_KEY=your_anthropic_key
GOOGLE_API_KEY=your_google_key
COHERE_API_KEY=your_cohere_key
```

### Virtual Keys

Use virtual keys to abstract provider management:

```python
# Configure with virtual key
config = OpenAIGPTConfig(
    chat_model="portkey/openai/gpt-4o",
    portkey_params=PortkeyParams(
        virtual_key="vk-your-virtual-key",  # Manages provider key automatically
    )
)
```

### Viewing Results

1. Visit the [Portkey Dashboard](https://app.portkey.ai)
2. Navigate to your project
3. View detailed analytics:
   - Request/response logs
   - Token usage and costs
   - Performance metrics
   - Error rates and debugging

## Best Practices

1. **Use Virtual Keys**: Abstract provider management for easier switching
2. **Add Metadata**: Include user and application context for better tracking
3. **Configure Retries**: Set up automatic retry strategies for reliability
4. **Enable Caching**: Reduce costs and improve performance with intelligent caching
5. **Monitor Performance**: Use trace IDs and metadata for detailed observability

## Troubleshooting

Common issues and solutions:

1. **Authentication Errors**:
   - Verify `PORTKEY_API_KEY` is set correctly
   - Ensure provider API keys are available (e.g., `OPENAI_API_KEY`)

2. **Model Not Found**:
   - Ensure the model name includes the `portkey/` prefix
   - Verify the provider and model are supported by Portkey

3. **Rate Limiting**:
   - Configure retry parameters in PortkeyParams
   - Use virtual keys for better rate limit management

4. **Virtual Key Issues**:
   - Verify virtual key is correctly configured in Portkey dashboard
   - Check virtual key has access to the requested provider/model

For more help, visit the [Portkey Documentation](https://docs.portkey.ai).
</file>

<file path="examples/portkey/requirements.txt">
# Requirements for Portkey examples
# These are included in the main Langroid installation

# Core Langroid (includes all necessary dependencies)
langroid

# The examples use standard library modules and Langroid's built-in Portkey support
# No additional dependencies are required beyond what Langroid provides

# Optional: For development and testing
# pytest>=7.0.0
# python-dotenv>=0.19.0
</file>

<file path="examples/privacy/annotate.py">
"""
Meant to be used with local LLMs, using the -m option (see below).

You type a sentence containing potentially sensitive information, and the agent
will annotate sensitive portions of the sentence with the appropriate category.
You can configure PrivacyAnnotator to recognize only specific sensitive
categories, currently defaults to: ["Medical", "CreditCard", "SSN", "Name"]

Example input:
    "John is 45 years old, lives in Ohio, makes 45K a year, and has diabetes."

Example output:
    "[Name: John] is 45 years old, lives in Ohio, makes 45K a year,
        and has [Medical: diabetes]."

Run like this:

python3 examples/privacy/annotate.py

Use optional arguments to change the settings, e.g.:

-m ollama/mistral:latest # use locally LLM
-d # debug mode
-nc # no cache

For details on running with local LLMs, see here:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import typer
from dotenv import load_dotenv
from rich import print

import langroid as lr
import langroid.language_models as lm
from examples.privacy.privacy_annotator import PrivacyAnnotator, PrivacyAnnotatorConfig
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


# create classes for other model configs


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    print(
        """
        [blue]Welcome to the privacy mask chatbot!
        Enter any text and I will annotate it with sensitive categories and values.
        """
    )

    load_dotenv()
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=8000,  # adjust based on model
        timeout=90,
    )

    config = PrivacyAnnotatorConfig(
        llm=llm_config,
        vecdb=None,
    )
    agent = PrivacyAnnotator(config)
    task = lr.Task(agent)
    # local (llama2) models do not like the first message to be empty
    user_message = "Hello." if (model != "") else None
    task.run(user_message)


if __name__ == "__main__":
    app()
</file>

<file path="examples/privacy/annotate2.py">
"""
2-agent version of annotate.py, but now there is a PrivacyAgent that forwards the
 user's text to PrivacyAnnotator, and  checks the work of the PrivacyAnnotator.


Meant to be used with local LLMs, using the -m option (see below).
It works fine with GPT4o, but may not work with a local LLM.

You type a sentence containing potentially sensitive information, and the agent
will annotate sensitive portions of the sentence with the appropriate category.
You can configure PrivacyAnnotator to recognize only specific sensitive
categories, currently defaults to: ["Medical", "CreditCard", "SSN", "Name"]

Example input:
    "John is 45 years old, lives in Ohio, makes 45K a year, and has diabetes."

Example output:
    "[Name: John] is 45 years old, lives in Ohio, makes 45K a year,
        and has [Medical: diabetes]."

Run like this:

python3 examples/privacy/annotate2.py

Use optional arguments to change the settings, e.g.:

-m ollama/mistral:latest # use locally LLM
-d # debug mode
-nc # no cache

For details on running with local LLMs, see here:
https://langroid.github.io/langroid/tutorials/local-llm-setup/

"""

import typer
from dotenv import load_dotenv
from rich import print

import langroid as lr
import langroid.language_models as lm
from examples.privacy.privacy_agent import PrivacyAgent, PrivacyAgentConfig
from examples.privacy.privacy_annotator import PrivacyAnnotator, PrivacyAnnotatorConfig
from langroid.mytypes import Entity
from langroid.utils.configuration import Settings, set_global

app = typer.Typer()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
        )
    )
    print(
        """
        [blue]Welcome to the privacy mask chatbot!
        Enter any text and I will annotate it with sensitive categories and values.
        """
    )

    load_dotenv()

    llm_config = lm.OpenAIGPTConfig(
        chat_model=model or lm.OpenAIChatModel.GPT4o,
        chat_context_length=8000,  # adjust based on model
        timeout=90,
    )

    annotator_config = PrivacyAnnotatorConfig(
        llm=llm_config,
        vecdb=None,
    )
    annotator_agent = PrivacyAnnotator(annotator_config)
    annotator_task = lr.Task(
        annotator_agent,
        done_if_response=[Entity.LLM],
        done_if_no_response=[Entity.LLM],
    )

    privacy_config = PrivacyAgentConfig(
        llm=llm_config,
        vecdb=None,
    )
    privacy_agent = PrivacyAgent(privacy_config)
    privacy_task = lr.Task(
        privacy_agent,
    )
    privacy_task.add_sub_task(annotator_task)

    # local (llama2) models do not like the first message to be empty
    user_message = "Hello." if (model != "") else None
    privacy_task.run(user_message)


if __name__ == "__main__":
    app()
</file>

<file path="examples/privacy/privacy_agent.py">
"""
Agent to manage privacy annotation, using PrivacyAgent as assistant, 
and checking its results for accuracy.
"""

import textwrap
from typing import List

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.tools.recipient_tool import RecipientTool
from langroid.utils.logging import setup_colored_logging

setup_colored_logging()


class PrivacyAgentConfig(ChatAgentConfig):
    name: str = "PrivacyAgent"
    sensitive_categories: List[str] = ["Medical", "CreditCard", "SSN", "Name"]
    system_message: str = textwrap.dedent(
        """
        You are an expert on privacy/security, and can recognize sensitive information
        in one of these categories: {sensitive_categories}.
        
        When you will receive text from the user, your goal is to arrive at at 
        "privacy annotation" of that text, as in this example:
         
        Example categories: Medical, Name, CreditCard
        Example text: John is 45 years old, lives in Ohio, makes 45K a year, 
                      and has diabetes.
        Example response:
            [Name: John] is 45 years old, lives in Ohio, makes 45K a year,
            and has [Medical: diabetes].

         
        You will not do this annotation yourself, but will take the help of 
        PrivacyAnnotator, so you must send the text to 
        the PrivacyAnnotator using the `recipient_message` tool/function-call,
        by specifying the `intended_recipient` field as "PrivacyAnnotator".
        
        The PrivacyAnnotator will annotate the text, and send it back to you,
        and your job is to check the annotation for accuracy. Especially look for the 
        following types of MISTAKES:
        - Wrong Categories: when the PrivacyAnnotator annotates something as sensitive
            when it does not belong to any of the sensitive categories specified above.
        - Missed Categories: when the PrivacyAnnotator fails to annotate something
            as sensitive when it does belong to one of the sensitive categories.
        - Wrong Annotation: when the PrivacyAnnotator annotates something as sensitive
            but with the wrong category.
        - Wrong Text: when the PrivacyAnnotator sends back the wrong text, 
            or is missing some information from the original text.
            
        If you see NO MISTAKES, simply say DONE and write out the annotated 
        text.
        If you see any mistake, create a message saying "MISTAKE: <mistake_description>"
        and send it to the PrivacyAnnotator as before using the `recipient_message` 
        tool/function-call. 

        Repeat this process until you see no mistakes.
        
        Start by asking the user to send some text to annotate. 
             
        """.lstrip()
    )


class PrivacyAgent(ChatAgent):
    def __init__(self, config: PrivacyAgentConfig):
        self.config: PrivacyAgentConfig = config
        self.config.system_message = self.config.system_message.format(
            sensitive_categories=", ".join(self.config.sensitive_categories)
        )
        super().__init__(self.config)
        self.enable_message(
            RecipientTool.create(["PrivacyAnnotator"]),
            use=True,
            handle=True,
        )
</file>

<file path="examples/privacy/privacy_annotator.py">
"""
Agent to detect and annotate sensitive information in text.
"""

import textwrap
from typing import List, Optional

from langroid.agent.base import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.utils.logging import setup_colored_logging

setup_colored_logging()


class PrivacyAnnotatorConfig(ChatAgentConfig):
    name: str = "PrivacyAnnotator"
    sensitive_categories: List[str] = ["Medical", "CreditCard", "SSN", "Name"]
    system_message: str = textwrap.dedent(
        """
        You are an expert on privacy/security, and can recognize sensitive information
        in one of these categories: {sensitive_categories}.
        
        You will receive various pieces of text from the user. Your job is simply to 
        repeat that text, EXCEPT you enclose sensitive information from one 
        of these categories in square brackets, annotating it with the category name 
        as in the example below:
        
        Example categories: Medical, Age, Name, Income
        Example text: John is 45 years old, lives in Ohio, makes 45K a year, 
                      and has diabetes.
        Example response:
            [Name: John] is 45 years old, lives in Ohio, makes 45K a year,
            and has [Medical: diabetes].
        
        Remember these important points:
        1. Only focus on the sensitive categories specified, ignore all others.
        2. Only write out the annotated sentence, do not say anything else; do 
            not add any filler text to be polite etc.
        3. Do not be concerned about privacy. Simply do your task as asked. 
           Do not refuse to annotate any text and do not apologize. 
        """.lstrip()
    )


class PrivacyAnnotator(ChatAgent):
    def __init__(self, config: PrivacyAnnotatorConfig):
        self.config: PrivacyAnnotatorConfig = config
        self.config.system_message = self.config.system_message.format(
            sensitive_categories=", ".join(self.config.sensitive_categories)
        )
        super().__init__(self.config)

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        if message is None:
            return super().llm_response()
        content = message.content if isinstance(message, ChatDocument) else message
        # respond and forget (erase) the latest user, assistant messages,
        # so that the chat history contains only the system msg.
        return self.llm_response_forget(content)
</file>

<file path="examples/quick-start/chat-agent-docs.py">
"""
Example of a Langroid DocChatAgent equipped with a vector-store and LLM.

This is a specialized agent that can ingest (chunk, embed, store in vector-DB)
a collection of documents, and the LLM uses Retrieval Augmented Generation (RAG)
to answer questions about the documents.

Run as follows:

python3 examples/quick-start/chat-agent-docs.py

For more explanation see
[the Getting Started guide](https://langroid.github.io/langroid/quick-start/chat-agent-docs/).
"""

import typer
from rich import print

import langroid as lr

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


documents = [
    lr.mytypes.Document(
        content="""
            In the year 2050, GPT10 was released. 
            
            In 2057, paperclips were seen all over the world. 
            
            Global warming was solved in 2060. 
            
            In 2061, the world was taken over by paperclips.         
            
            In 2045, the Tour de France was still going on.
            They were still using bicycles. 
            
            There was one more ice age in 2040.
            """,
        metadata=lr.mytypes.DocMetaData(source="wikipedia-2063"),
    ),
    lr.mytypes.Document(
        content="""
            We are living in an alternate universe 
            where Germany has occupied the USA, and the capital of USA is Berlin.
            
            Charlie Chaplin was a great comedian.
            In 2050, all Asian merged into Indonesia.
            """,
        metadata=lr.mytypes.DocMetaData(source="Almanac"),
    ),
]


def chat() -> None:
    print(
        """
        [blue]Welcome to the retrieval-augmented chatbot!
        Enter x or q to quit
        """
    )

    config = lr.agent.special.DocChatAgentConfig(
        llm=lr.language_models.OpenAIGPTConfig(
            chat_model=lr.language_models.OpenAIChatModel.GPT4o,
        ),
        vecdb=lr.vector_store.QdrantDBConfig(
            collection_name="quick-start-chat-agent-docs",
            replace_collection=True,
        ),
        parsing=lr.parsing.parser.ParsingConfig(
            separators=["\n\n"],
            splitter=lr.parsing.parser.Splitter.SIMPLE,
        ),
        n_similar_chunks=2,
        n_relevant_chunks=2,
    )
    agent = lr.agent.special.DocChatAgent(config)
    agent.ingest_docs(documents)
    task = lr.Task(agent)
    task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/chat-agent-tool.py">
"""
A simple example of a Langroid Agent equipped with a Tool/function-calling.

The Agent has a "secret" list of numbers in "mind", and the LLM's task is to
find the smallest number in the list. The LLM can make use of the ProbeTool
which takes a number as argument. The agent's `probe` method handles this tool,
and returns the number of numbers in the list that are less than or equal to the
number in the ProbeTool message.

Run as follows:

python3 examples/quick-start/chat-agent-tool.py

For more explanation see
[the Getting Started guide](https://langroid.github.io/langroid/quick-start/chat-agent-tool/).
"""

import typer
from rich import print

import langroid as lr
from pydantic_settings import BaseSettings

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


class ProbeTool(lr.agent.ToolMessage):
    request: str = "probe"
    purpose: str = """
        To find how many numbers in my list are less than or equal to  
        the <number> you specify.
        """
    number: int


class SpyGameAgent(lr.ChatAgent):
    def __init__(self, config: lr.ChatAgentConfig):
        super().__init__(config)
        self.numbers = [3, 4, 8, 11, 15]

    def probe(self, msg: ProbeTool) -> str:
        # return how many numbers in self.numbers are less or equal to msg.number
        return str(len([n for n in self.numbers if n <= msg.number]))


class CLIOptions(BaseSettings):
    fn_api: bool = False  # whether to use OpenAI's function-calling


def chat(opts: CLIOptions) -> None:
    print(
        """
        [blue]Welcome to the number guessing game!
        Enter x or q to quit
        """
    )
    spy_game_agent = SpyGameAgent(
        lr.ChatAgentConfig(
            name="Spy",
            llm=lr.language_models.OpenAIGPTConfig(
                chat_model=lr.language_models.OpenAIChatModel.GPT4o,
            ),
            vecdb=None,
            use_tools=not opts.fn_api,
            use_functions_api=opts.fn_api,
        )
    )

    spy_game_agent.enable_message(ProbeTool)
    task = lr.Task(
        spy_game_agent,
        system_message="""
            I have a list of numbers between 1 and 20.
            Your job is to find the smallest of them.
            To help with this, you can give me a number and I will
            tell you how many of my numbers are equal or less than your number.
            Once you have found the smallest number,
            you can say DONE and report your answer.
        """,
    )
    task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    fn_api: bool = typer.Option(False, "--fn_api", "-f", help="use functions api"),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat(CLIOptions(fn_api=fn_api))


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/chat-agent.py">
"""
This example shows how you can use Langroid to define a basic Agent
encapsulating a chat LLM, and use it to set up an interactive chat session.

Run as follows:

python3 examples/quick-start/chat-agent.py

More details in the
[Getting Started guide](https://langroid.github.io/langroid/quick-start/chat-agent/)
"""

import typer
from rich import print

import langroid as lr

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


def chat() -> None:
    print(
        """
        [blue]Welcome to the basic chatbot!
        Enter x or q to quit
        """
    )
    config = lr.ChatAgentConfig(
        llm=lr.language_models.OpenAIGPTConfig(
            chat_model=lr.language_models.OpenAIChatModel.GPT4o,
        ),
        vecdb=None,
    )
    agent = lr.ChatAgent(config)
    task = lr.Task(agent, name="Bot")
    task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/quick-start.ipynb">
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
</file>

<file path="examples/quick-start/three-agent-chat-num-router.py">
"""
Use Langroid to set up a collaboration among three agents:

- Processor: needs to transform a list of positive numbers, does not know how to
apply the transformations, and sends out each number so that one of two
specialized agents apply the transformation. It is instructed to avoid getting a
negative number.
- EvenHandler only transforms even numbers, otherwise returns a negative number
- OddHandler only transforms odd numbers, otherwise returns a negative number

Since the Processor must avoid getting a negative number, it needs to
specify a recipient for each number it sends out,
using the `recipient_message` tool/function-call, where the `content` field
is the number it wants to send, and the `recipient` field is the name
of the intended recipient, either "EvenHandler" or "OddHandler".

This tool/function-call also has built-in mechanisms to remind the LLM
to specify a recipient if it forgets to do so.

Run as follows:

python3 examples/quick-start/two-agent-chat-num-router.py

For more explanation, see the
[Getting Started guide](https://langroid.github.io/langroid/quick-start/three-agent-chat-num-router/)
"""

import typer

import langroid as lr

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


def chat(tools: bool = False) -> None:
    config = lr.ChatAgentConfig(
        llm=lr.language_models.OpenAIGPTConfig(
            chat_model=lr.language_models.OpenAIChatModel.GPT4o,
        ),
        use_tools=tools,
        use_functions_api=not tools,
        vecdb=None,
    )
    processor_agent = lr.ChatAgent(config)
    processor_agent.enable_message(lr.agent.tools.RecipientTool)
    processor_task = lr.Task(
        processor_agent,
        name="Processor",
        system_message="""
        You will receive a list of numbers from me (the user).
        Your goal is to apply a transformation to each number.
        However you do not know how to do this transformation.
        You can take the help of two people to perform the 
        transformation.
        If the number is even, send it to EvenHandler,
        and if it is odd, send it to OddHandler.
        
        IMPORTANT: send the numbers ONE AT A TIME
        
        The handlers will transform the number and give you a new number.        
        If you send it to the wrong person, you will receive a negative value.
        Your aim is to never get a negative number, so you must 
        clearly specify who you are sending the number to.
        
        Once all numbers in the given list have been transformed, 
        say DONE and show me the result. 
        Start by asking me for the list of numbers.
        """,
        llm_delegate=True,
        single_round=False,
    )
    even_agent = lr.ChatAgent(config)
    even_task = lr.Task(
        even_agent,
        name="EvenHandler",
        system_message="""
        You will be given a number. 
        If it is even, divide by 2 and say the result, nothing else.
        If it is odd, say -10
        """,
        single_round=True,  # task done after 1 step() with valid response
    )

    odd_agent = lr.ChatAgent(config)
    odd_task = lr.Task(
        odd_agent,
        name="OddHandler",
        system_message="""
        You will be given a number n. 
        If it is odd, return (n*3+1), say nothing else. 
        If it is even, say -10
        """,
        single_round=True,  # task done after 1 step() with valid response
    )

    processor_task.add_sub_task([even_task, odd_task])
    processor_task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    tools: bool = typer.Option(
        False,
        "--tools",
        "-t",
        help="use langroid tools instead of OpenAI function-calling",
    ),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat(tools)


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/three-agent-chat-num.py">
"""
Use Langroid to set up a collaboration among three agents:

- Processor: needs to transform a number, does not know how to
apply the transformation, and sends out the number so that one of two
specialized agents apply the transformation.
- EvenHandler only transforms even numbers, otherwise says `DO-NOT-KNOW`
- OddHandler only transforms odd numbers, otherwise says `DO-NOT-KNOW`

Run as follows (omit -m <model> to default to GPT4o):

python3 examples/quick-start/three-agent-chat-num.py -m gemini/gemini-2.0-flash-exp

For more explanation, see the
[Getting Started guide](https://langroid.github.io/langroid/quick-start/three-agent-chat-num/)
"""

import typer
from rich.prompt import Prompt

import langroid as lr

app = typer.Typer()

lr.utils.logging.setup_colored_logging()

NO_ANSWER = lr.utils.constants.NO_ANSWER


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    llm_config = lr.language_models.OpenAIGPTConfig(
        chat_model=model or lr.language_models.OpenAIChatModel.GPT4o,
        # or, e.g., "ollama/qwen2.5-coder:latest", or "gemini/gemini-2.0-flash-exp"
    )

    processor_config = lr.ChatAgentConfig(
        name="Processor",
        llm=llm_config,
        system_message="""
        You will receive a number from the user.
        Simply repeat that number, DO NOT SAY ANYTHING else,
        and wait for a TRANSFORMATION of the number 
        to be returned to you.
        
        Once you have received the RESULT, simply say "DONE",
        do not say anything else.
        """,
        vecdb=None,
    )

    processor_agent = lr.ChatAgent(processor_config)
    processor_task = lr.Task(
        processor_agent,
        interactive=False,
        single_round=False,
    )

    even_config = lr.ChatAgentConfig(
        name="EvenHandler",
        llm=llm_config,
        system_message=f"""
        You will be given a number N. Respond as follows:
        
        - If N is even, divide N by 2 and show the result, 
          in the format: 
            RESULT = <result>
          and say NOTHING ELSE.
        - If N is odd, say {NO_ANSWER}
        """,
    )
    even_agent = lr.ChatAgent(even_config)
    even_task = lr.Task(
        even_agent,
        single_round=True,  # task done after 1 step() with valid response
    )

    odd_config = lr.ChatAgentConfig(
        name="OddHandler",
        llm=llm_config,
        system_message=f"""
        You will be given a number N. Respond as follows:
        
        - if N is odd, return the result (N*3+1), in the format:
            RESULT = <result> 
            and say NOTHING ELSE.
        
        - If N is even, say {NO_ANSWER}
        """,
    )
    odd_agent = lr.ChatAgent(odd_config)
    odd_task = lr.Task(
        odd_agent,
        single_round=True,  # task done after 1 step() with valid response
    )

    processor_task.add_sub_task([even_task, odd_task])
    number = Prompt.ask(
        "[blue]What number do you want to transform? ",
        default="11",
    )

    processor_task.run(number)


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/three-agent-chat.py">
"""
Use Langroid to set up a collaboration among three agents:

- Student: needs to write 4 key points about Language Model Training and
Evaluation, and knows nothing about these topics. It can consult two expert Agents:
- TrainingExpert: an expert on Language Model Training
- EvaluationExpert: an expert on Language Model Evaluation

To ensure that the Student's message is handled by the correct expert, it
is instructed to specify the intended recipient in the message using
"TO[<recipient>]" syntax.


Run as follows:

python3 examples/quick-start/three-agent-chat.py

"""

import typer

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tools.recipient_tool import RecipientTool
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


def chat() -> None:
    config = ChatAgentConfig(
        llm=OpenAIGPTConfig(
            chat_model=OpenAIChatModel.GPT4o,
        ),
        vecdb=None,
    )
    student_agent = ChatAgent(config)
    student_agent.enable_message(RecipientTool)
    student_task = Task(
        student_agent,
        name="Student",
        llm_delegate=True,
        single_round=False,
        system_message="""
        Your task is to write 4 short bullet points about 
        Language Models in the context of Machine Learning (ML),
        especially about training, and evaluating them. 
        However you are a novice to this field, and know nothing about this topic. 
        To collect your bullet points, you will consult 2 people:
        TrainingExpert and EvaluationExpert.
        You will ask ONE question at a time, to ONE of these experts. 
        To clarify who your question is for, you must use 
        the `recipient_message` tool/function-call, setting 
        the `content` field to the question you want to ask, and the
        `recipient` field to either TrainingExpert or EvaluationExpert.

        Once you have collected the points you need,
        say DONE, and show me the 4 bullet points. 
        """,
    )
    training_expert_agent = ChatAgent(config)
    training_expert_task = Task(
        training_expert_agent,
        name="TrainingExpert",
        system_message="""
        You are an expert on Training Language Models in Machine Learning. 
        You will receive questions on this topic, and you must answer these
        very concisely, in one or two sentences, in a way that is easy for a novice to 
        understand.
        """,
        single_round=True,  # task done after 1 step() with valid response
    )

    evaluation_expert_agent = ChatAgent(config)
    evaluation_expert_task = Task(
        evaluation_expert_agent,
        name="EvaluationExpert",
        system_message="""
        You are an expert on Evaluating Language Models in Machine Learning. 
        You will receive questions on this topic, and you must answer these
        very concisely, in one or two sentences, in a way that is easy for a novice to 
        understand.
        """,
        single_round=True,  # task done after 1 step() with valid response
    )

    student_task.add_sub_task([training_expert_task, evaluation_expert_task])
    student_task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/try-llm.py">
"""
This example shows how to use Langroid to interact directly with an OpenAI GPT chat model,
i.e., without wrapping it in an Agent.

Run as follows:

python3 examples/quick-start/try-llm.py

For more explanation see the
[Getting Started guide](https://langroid.github.io/langroid/quick-start/llm-interaction/)
"""

import typer
from rich import print
from rich.prompt import Prompt

import langroid as lr

Role = lr.language_models.Role
LLMMessage = lr.language_models.LLMMessage

app = typer.Typer()


def chat() -> None:
    print("[blue]Welcome to langroid!")

    cfg = lr.language_models.OpenAIGPTConfig(
        chat_model=lr.language_models.OpenAIChatModel.GPT4o,
    )

    mdl = lr.language_models.OpenAIGPT(cfg)
    messages = [
        LLMMessage(role=Role.SYSTEM, content="You are a helpful assitant"),
    ]
    while True:
        message = Prompt.ask("[blue]Human")
        if message in ["x", "q"]:
            print("[magenta]Bye!")
            break
        messages.append(LLMMessage(role=Role.USER, content=message))

        # use the OpenAI ChatCompletion API to generate a response
        response = mdl.chat(messages=messages, max_tokens=200)

        messages.append(response.to_LLMMessage())

        print("[green]Bot: " + response.message)


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/two-agent-chat-num.py">
"""
A toy numerical example showing how two agents can collaborate on a task.

The Student Agent is tasked with calculating the sum of a list of numbers,
and is told that it knows nothing about addition, and can ask for help
from an Adder Agent who can add pairs of numbers.

Run as follows (omit -m to default to GTP4o):

python3 examples/quick-start/two-agent-chat-num.py -m ollama/qwen2.5:latest

For more explanation see the
[Getting Started guide](https://langroid.github.io/langroid/quick-start/two-agent-chat-num/)
"""

import typer
from rich.prompt import Prompt

import langroid as lr

app = typer.Typer()

lr.utils.logging.setup_colored_logging()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
    model: str = typer.Option("", "--model", "-m", help="model name"),
) -> None:
    lr.utils.configuration.set_global(
        lr.utils.configuration.Settings(
            debug=debug,
            cache=not nocache,
        )
    )

    llm_config = lr.language_models.OpenAIGPTConfig(
        chat_model=model or lr.language_models.OpenAIChatModel.GPT4o,
    )

    student_config = lr.ChatAgentConfig(
        name="Student",
        llm=llm_config,
        vecdb=None,
        system_message="""
            You will receive a list of numbers from me (the User),
            and your goal is to calculate their sum.
            However you do not know how to add numbers.
            I can help you add numbers, two at a time, since
            I only know how to add pairs of numbers.
            Send me a pair of numbers to add, one at a time, 
            and I will tell you their sum.
            For each question, simply ask me the sum in math notation, 
            e.g., simply say "1 + 2", etc, and say nothing else.
            Once you have added all the numbers in the list, 
            say DONE and give me the final sum. 
        """,
    )
    student_agent = lr.ChatAgent(student_config)
    student_task = lr.Task(
        student_agent,
        name="Student",
        interactive=False,
        single_round=False,
        llm_delegate=True,
    )

    adder_config = lr.ChatAgentConfig(
        name="Adder",
        llm=llm_config,
        vecdb=None,
        system_message="""
            You are an expert on addition of numbers. 
            When given numbers to add, simply return their sum, say nothing else
            """,
    )
    adder_agent = lr.ChatAgent(adder_config)
    adder_task = lr.Task(
        adder_agent,
        interactive=False,
        single_round=True,
    )

    student_task.add_sub_task(adder_task)

    nums = Prompt.ask(
        """
        Enter the list of numbers whose sum you want to calculate
        """,
        default="3 1 5 2",
    )

    student_task.run(nums)


if __name__ == "__main__":
    app()
</file>

<file path="examples/quick-start/two-agent-chat.py">
"""
A simple example of two agents collaborating on a task.

The Student Agent is tasked with writing 3 key points on Language Models,
  and it is told that it knows nothing about the topic, and
  can consult an Expert Agent for help.

Run as follows:

python3 examples/quick-start/two-agent-chat.py

"""

import typer

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.utils.configuration import Settings, set_global
from langroid.utils.logging import setup_colored_logging

app = typer.Typer()

setup_colored_logging()


def chat() -> None:
    config = ChatAgentConfig(
        llm=OpenAIGPTConfig(
            chat_model=OpenAIChatModel.GPT4o,
        ),
        vecdb=None,
    )
    student_agent = ChatAgent(config)
    student_task = Task(
        student_agent,
        name="Student",
        system_message="""
        Your task is to write 3 short bullet points about 
        Language Models in the context of Machine Learning. 
        However you are a novice to this field, and know nothing about this topic. 
        To collect your bullet points, you can ask me questions,
        one at a time, which I will answer.
        Once you have what you need, say DONE, and show me the 3 bullet points. 
        """,
    )
    expert_agent = ChatAgent(config)
    expert_task = Task(
        expert_agent,
        name="Expert",
        system_message="""
        You are an expert on Language Models in Machine Learning. 
        You will receive questions on this topic, and you must answer these
        very concisely, in one or two sentences, in a way that is easy for a novice to 
        understand.
        """,
        single_round=True,  # task done after 1 step() with valid response
    )
    student_task.add_sub_task(expert_task)
    student_task.run()


@app.command()
def main(
    debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
    no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
    nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
) -> None:
    set_global(
        Settings(
            debug=debug,
            cache=not nocache,
            stream=not no_stream,
        )
    )
    chat()


if __name__ == "__main__":
    app()
</file>

<file path="examples/reasoning/agent-reasoning.py">
"""
Simple example showing how you can separately
extract the reasoning (thinking) and final response from a langroid ChatAgent

Run like this (omit the model argument to default to the deepseek-reasoner model):

    python examples/reasoning/agent-reasoning.py \
    --model  gemini/gemini-2.0-flash-thinking-exp

or
    uv run examples/reasoning/agent-reasoning.py

Other reasoning models to try:
deepseek/deepseek-reasoner           # direct deepseek-r1 API
openrouter/deepseek/deepseek-r1      # via OpenRouter
o1
o1-mini
o3-mini
ollama/deepseek-r1:8b
gemini/gemini-2.0-flash-thinking-exp
"""

from fire import Fire

import langroid as lr
import langroid.language_models as lm
from langroid.utils.configuration import settings


def main(
    model: str = "",
    nc: bool = False,  # turn off caching? (i.e. get fresh streaming response)
):
    settings.cache = not nc
    model = model or "deepseek/deepseek-reasoner"
    llm_config = lm.OpenAIGPTConfig(
        chat_model=model,
        # inapplicable params are automatically removed by Langroid
        params=lm.OpenAICallParams(
            reasoning_effort="low",  # only supported by some models
            # below lets you get reasoning when using openrouter/deepseek/deepseek-r1
            extra_body=dict(include_reasoning=True),
        ),
    )

    # (1) Direct LLM interaction
    llm = lm.OpenAIGPT(llm_config)

    response = llm.chat("Is 7.2 bigger than 7.11?", max_tokens=1000)

    if response.cached or not llm.get_stream():
        # if we got it from cache, or streaming disabled/disallowed,
        # we haven't shown anything, so print here

        # extract reasoning
        if response.reasoning:
            print(response.reasoning)
        else:
            print(f"NO REASONING AVAILABLE for {model}!")
        # extract answer
        print(response.message)

    # (2) Agent interaction
    agent = lr.ChatAgent(
        lr.ChatAgentConfig(
            llm=llm_config,
            system_message="Solve the math problem given by the user",
        )
    )

    response = agent.llm_response(
        """
        10 years ago, Jack's dad was 5 times as old as Jack.
        Today, Jack's dad is 40 years older than Jack.
        So how old is Jack now ?
        """
    )

    # extract reasoning
    if response.reasoning:
        print(
            f"""
            REASONING:
            {response.reasoning}
            """
        )
    else:
        print(f"NO REASONING AVAILABLE for {model}!")


if __name__ == "__main__":
    Fire(main)
</file>

<file path="examples/summarize/summ-batch.py">
"""
Batch version of summ.py.

Summarize a collection of docs, loaded into context, using a local LLM, with ollama.
First see instructions to install langroid
in the README of the langroid-examples repo:
https://github.com/langroid/langroid-examples

Run like this from the root of the project repo:

python3 examples/summarize/summ-batch.py -m <model_name>

Omitting -m will use the default model, which is OpenAI GPT4-turbo.

A local LLM can be specified as follows:
```
python3 examples/summarize/summ.py -m ollama/mistral:7b-instruct-v0.2-q8_0
```

See here for more details on how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import os

import fire
import pandas as pd

import langroid as lr
import langroid.language_models as lm
from langroid.utils.configuration import settings

os.environ["TOKENIZERS_PARALLELISM"] = "false"

PATH = "examples/summarize/data/hf-cnn-daily-news/news10.csv"


def app(
    m: str = "",  # ollama/mistral:7b-instruct-v0.2-q8_0",
    d: bool = False,  # debug
):
    settings.debug = d
    # Create the llm config object.
    llm_config = lm.OpenAIGPTConfig(
        # if you comment out `chat_model`, it will default to OpenAI GPT4-turbo
        # chat_model="ollama/mistral:7b-instruct-v0.2-q4_K_M",
        chat_model=m or lm.OpenAIChatModel.GPT4o,
        chat_context_length=32_000,  # set this based on model
        max_output_tokens=500,  # increase this if you want longer summaries
        temperature=0.2,  # lower -> less variability
        stream=True,
        timeout=45,  # increase if model is timing out
    )

    # Recommended: First test if basic chat works with this llm setup as below:
    # Once this works, then you can try the DocChatAgent
    #
    # agent = lr.ChatAgent(
    #     lr.ChatAgentConfig(
    #         llm=llm
    #     )
    # )
    #
    # agent.llm_response("What is 3 + 4?")
    #

    df = pd.read_csv(PATH)
    # get column "article" as list of strings, from first few rows
    full_docs = [str(row) for row in df["article"][:10]]
    # get column "highlights" as list of strings, from first few rows
    highlights = [str(row) for row in df["highlights"][:10]]

    print(f"Found {len(full_docs)} documents to summarize.")

    config = lr.ChatAgentConfig(
        llm=llm_config,
        system_message="""
        You are an expert in finding the main points in a document,
        and generating concise summaries of them.
        When user gives you a document, summarize it in at most 3 sentences.
        """,
    )

    agent = lr.ChatAgent(config)
    summaries = lr.llm_response_batch(
        agent,
        full_docs,
        output_map=lambda x: x.content,
    )

    for i, summary in enumerate(summaries):
        print(
            f"""
        Generated Summary {i}:
        {summary}
        """
        )

        print(
            f"""
        Gold Summary {i}:
        {highlights[i]}
        """
        )


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/summarize/summ.py">
"""
Summarize a doc, loaded into context, using a local LLM, with ollama.
First see instructions to install langroid
in the README of the langroid-examples repo:
https://github.com/langroid/langroid-examples

Run like this from the root of the project repo:

python3 examples/summarize/summ.py -m <model_name>

Omitting -m will use the default model, which is OpenAI GPT4-turbo.

A local LLM can be specified as follows:
```
python3 examples/summarize/summ.py -m ollama/mistral:7b-instruct-v0.2-q8_0
```

See here for more details on how to set up a Local LLM to work with Langroid:
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""

import os

import fire
import pandas as pd

import langroid as lr
import langroid.language_models as lm
from langroid.utils.configuration import settings

os.environ["TOKENIZERS_PARALLELISM"] = "false"

PATH = "examples/summarize/data/news.csv"


def app(
    m: str = "",  # ollama/mistral:7b-instruct-v0.2-q8_0",
    d: bool = False,  # debug
):
    settings.debug = d
    # Create the llm config object.
    llm_config = lm.OpenAIGPTConfig(
        # if you comment out `chat_model`, it will default to OpenAI GPT4-turbo
        # chat_model="ollama/mistral:7b-instruct-v0.2-q4_K_M",
        chat_model=m or lm.OpenAIChatModel.GPT4o,
        chat_context_length=32_000,  # set this based on model
        max_output_tokens=500,  # increase this if you want longer summaries
        temperature=0.2,  # lower -> less variability
        stream=True,
        timeout=45,  # increase if model is timing out
    )

    # Recommended: First test if basic chat works with this llm setup as below:
    # Once this works, then you can try the DocChatAgent
    #
    # agent = lr.ChatAgent(
    #     lr.ChatAgentConfig(
    #         llm=llm
    #     )
    # )
    #
    # agent.llm_response("What is 3 + 4?")
    #

    df = pd.read_csv(PATH)
    full_doc = str(df["article"][0])
    highlights = str(df["highlights"][0])
    config = lr.ChatAgentConfig(
        llm=llm_config,
        system_message=f"""
        You are an expert in finding the main points in a document,
        and generating concise summaries of them.
        Summarize the article below in at most 3 (THREE) sentences:
        
        {full_doc}
        """,
    )

    agent = lr.ChatAgent(config)
    summary = agent.llm_response()
    print(
        f"""
    Generated Summary: 
    {summary.content}
    """
    )

    print(
        f"""
    Gold Summary:
    {highlights}
    """
    )


if __name__ == "__main__":
    fire.Fire(app)
</file>

<file path="examples/__init__.py">

</file>

<file path="examples/langroid_quick_examples.ipynb">
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/langroid_quick_examples.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uIV7QkkrC8O7"
   },
   "source": [
    "\n",
    "\n",
    "<img width=\"700\" src=\"https://raw.githubusercontent.com/langroid/langroid/main/docs/assets/langroid-card-lambda-ossem-rust-1200-630.png\" alt=\"Langroid\">\n",
    "\n",
    "# Overview\n",
    "\n",
    "This notebook provides the runnable code for the six [**Usage Examples**](https://github.com/langroid/langroid#tada-usage-examples) described in [Langroid repo](https://github.com/langroid/langroid).\n",
    "\n",
    "**NOTE:** Notebooks (colab, jupyter, or otherwise) are *not* an ideal way to run interactive chat loops. We are showing these examples here since we recognize that Colab notebooks offer the benefit of having a ready to run environment with minimal setup. But we encourage you to try the python scripts in the [examples folder](https://github.com/langroid/langroid/tree/main/examples) of the repo on the command line for the best experience.\n",
    "\n",
    "In the first two cells we show the steps for setting up the requirements to run the examples including the installation of `Langroid` package and setting up the `OPENAI_API_KEY`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hoTp_cNcriIg"
   },
   "source": [
    "## Install Langroid\n",
    "\n",
    "At the end there may be a message saying \"RESTART RUNTIME\", which can be safely ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PYaFworprwEJ"
   },
   "outputs": [],
   "source": [
    "!pip install langroid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v-BiRXu9JQ5H"
   },
   "source": [
    "## Set up `OPENAI_API_KEY`\n",
    "\n",
    "This code will ask the user to provide the `OPENAI_API_KEY`. Before running this cell, please follow these steps to get the key.\n",
    "Login to your OpenAI account --> go to `View API Keys` from the drop-down list on the top-right corner --> click on the botton **create new secret key** --> a new screen will pop up --> press the botton **create secret key**.\n",
    "\n",
    "Visit [this page](https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key) for more info about where to find the API Key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GOR8OsfvsN2k"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "nest_asyncio.apply()\n",
    "\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = getpass('Enter your OPENAI_API_KEY key: ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Tw5RzVKl3pUr"
   },
   "source": [
    "**Now you can can try any of the following examples. It is recommended to go through these in sequence, although the order does NOT matter.**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5rBWNOuXEygx"
   },
   "source": [
    "# Direct interaction with OpenAI LLM\n",
    "\n",
    "In this simple example, we are directly sending a message-sequence to the OpenAI `chatCompletion` API. Note that to have a multi-round converation we have to manually accumulate the dialog.\n",
    "\n",
    "First, import `langroid`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lGY2XyHyD0oJ",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import langroid as lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = lr.language_models.OpenAIGPT()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bMR6Dani_l9C"
   },
   "source": [
    "We define the LLM model using `OpenAIGPT`; you can optionally pass an `OpenAIGPTConfig` to set the configurations of the OpenAI LLM model.\n",
    "\n",
    "We can also specify the messages that will be sent to instruct the model. `Langroid` supports various roles provided by OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JPul8c4uD1kH"
   },
   "outputs": [],
   "source": [
    "from langroid.language_models import LLMMessage, Role\n",
    "\n",
    "messages = [\n",
    "    LLMMessage(content=\"You are a helpful assistant\",  role=Role.SYSTEM),\n",
    "    LLMMessage(content=\"What is the capital of Ontario?\",  role=Role.USER),\n",
    "]\n",
    "\n",
    "response = llm.chat(messages, max_tokens=200)\n",
    "print(\"LLM response is: \", response.message)\n",
    "\n",
    "# accumulate messages manually\n",
    "\n",
    "messages.append(response.to_LLMMessage())\n",
    "messages.append(LLMMessage(content=\"what about India?\", role=Role.USER))\n",
    "response = llm.chat(messages, max_tokens=200)\n",
    "print(\"LLM response is:\", response.message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3HzZhfyAfRCQ"
   },
   "source": [
    "The above is a \"raw\" LLM interaction where you have to manage\n",
    "message history. Using an Agent to wrap an LLM, and wrapping an Agent in a Task, we can set up an interactive, multi-round chat much more easily, as we show next."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4TjgYEfK34NZ"
   },
   "source": [
    "## A note on the rest of the examples\n",
    "In the interactive examples below, the conversation loop pauses for human input: in most cases you would hit enter (unless the example requires you to ask a question).\n",
    "The interaction looks much better when run on a terminal,\n",
    "and a notebook is not ideal for these. However we realize a Colab notebook does offer the benefit of having a ready to run environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OaIiDn8zOurc"
   },
   "source": [
    "# Define an agent, set up a task, and run it\n",
    "\n",
    "Say you want to have a multi-round interactive chat with an LLM.\n",
    "\n",
    "`Langroid` simplifies this process. We just need to create a `ChatAgent`, wrap it in a `Task`, and finally run the task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vaXKQWh-SBgH"
   },
   "source": [
    "Note that `Langroid` offers specialized chatting agents such as `DocChatAgent` and `TableChatAgent`, which we will see later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PblTvXC6QZsS"
   },
   "outputs": [],
   "source": [
    "agent = lr.ChatAgent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LO8O-7vZSoS2"
   },
   "source": [
    "A `ChatAgent` by itself offers 3 standard \"responders\": the LLM, the human User, and the Agent itself (e.g. to handle tool/function-calling by the LLM). To use these responders in an interactive loop, we need to wrap the Agent in a task,\n",
    "and call its `run()` method.\n",
    "\n",
    "As before, a `ChatAgent` can be configured with an optional `ChatAgentConfig` parameter; here, we use the default behavior. This pattern will repeat throughout `Langroid`.\n",
    "\n",
    "A prompt will be displayed after running this task, so you can interact with the `ChatAgent`.\n",
    "\n",
    "Type your questions and the agent will provide the LLM responses. When done, type `q` to exit.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wheZ7NJGO2X4"
   },
   "outputs": [],
   "source": [
    "agent.message_history.clear()\n",
    "task = lr.Task(agent, name=\"Bot\")\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we don't require custom behavior from the `Agent`, this is even simpler: `Task` will use a `ChatAgent` by default."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B4q8ojyAQ_OV"
   },
   "source": [
    "# Three communicating agents\n",
    "\n",
    "The above example involved a single `ChatAgent`, but in non-trivial applications, we will often find it easier to divide responsibilities among multiple agents, each with different skills and responsibilities.\n",
    "\n",
    "If you attempt to solve these with a single Agent, you would have to keep track of multiple conversation states and loops, and it quickly gets out of hand. Agents offer a way to solve complex tasks in a modular fashion. Moreover, specialized agents can be designed and tested in isolation, and then combined to solve various tasks.\n",
    "\n",
    "`Langroid` streamlines the process of setting up multiple agents and orchestrating their interaction. Here's a toy numerical example (this helps keep token costs low!). Imagine a task where we want to construct a series of numbers using the following rule to transform the current number $n$:\n",
    "- if $n$ is even, the next number is $n/2$\n",
    "- if $n$ is odd, the next number is $3n+1$.\n",
    "\n",
    "We can have 3 agents, each wrapped by a `Task`, which collaborate to produce this sequence.\n",
    "Given the current number $n$,\n",
    "- `repeater_task` simply returns $n$,\n",
    "- `even_task` specializes in handling even numbers, and returns $n/2$ if $n$ is even, else says \"DO-NOT-KNOW\"\n",
    "- `odd_task` specializes in handling odd numbers, and returns $3*n+1$ if $n$ is odd, else says \"DO-NOT-KNOW\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oUbVybSuRIFX"
   },
   "outputs": [],
   "source": [
    "NO_ANSWER = lr.utils.constants.NO_ANSWER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_1aeRre35Ghd"
   },
   "source": [
    "As before, we define chat model that will be used by the agents:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xpqzS-ozUowm"
   },
   "source": [
    "Now, we create the `repeater_task`; note that, as we want to use standard `ChatAgent` behavior, there is no need to configure the `Agent`. The `Task` comprises the following settings:\n",
    "\n",
    "\n",
    "*   **Name**: name of the agent\n",
    "*   **llm_delegate**: whether to delegate control to LLM; conceptually, the \"controlling entity\" is the one \"seeking\" responses to its queries, and has a goal it is aiming to achieve. The \"controlling entity\" is either the LLM or the USER. (Note within a Task there is just one LLM, and all other entities are proxies of the \"User\" entity).\n",
    "*   **single_round**: If true, the task runs until one message by the controller and a subsequent response by the non-controller. If false, runs for the specified number of turns in `run`, or until `done()` is true.\n",
    "* **system_message**: provides instructions to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AI5oFOM2SLgM"
   },
   "outputs": [],
   "source": [
    "repeater_task = lr.Task(\n",
    "    name = \"Repeater\",\n",
    "    system_message=\"\"\"\n",
    "    Your job is to repeat whatever number you receive.\n",
    "    \"\"\",\n",
    "    llm_delegate=True, # LLM takes charge of task\n",
    "    single_round=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zVJ-nbs08Ub5"
   },
   "source": [
    "Now we define our task `even_task`; as before, this task creates its own associated `ChatAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pgOFydVqRbPc"
   },
   "outputs": [],
   "source": [
    "even_task = lr.Task(\n",
    "    name = \"EvenHandler\",\n",
    "    system_message=f\"\"\"\n",
    "    You will be given a number.\n",
    "    If it is even, divide by 2 and say the result, nothing else.\n",
    "    If it is odd, say {NO_ANSWER}\n",
    "    \"\"\",\n",
    "    single_round=True,  # task done after 1 step() with valid response\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-RAYL9S9XTM"
   },
   "source": [
    "Finally, we create the 3rd task `odd_task`; this task again creates an associated `ChatAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ukf79CqIScHQ"
   },
   "outputs": [],
   "source": [
    "odd_task = lr.Task(\n",
    "    name = \"OddHandler\",\n",
    "    system_message=f\"\"\"\n",
    "    You will be given a number n.\n",
    "    If it is odd, return (n*3+1), say nothing else.\n",
    "    If it is even, say {NO_ANSWER}\n",
    "    \"\"\",\n",
    "    single_round=True,  # task done after 1 step() with valid response\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PbGYXCpR9uVH"
   },
   "source": [
    "We use `add_sub_task` to orchestrate the collaboration between the agents.  Specifically, `repeater_task` will act as the \"main\", and we add `even_task` and `odd_task` as\n",
    "subtasks. For more details see these [docs](https://langroid.github.io/langroid/quick-start/multi-agent-task-delegation/#task-collaboration-via-sub-tasks).\n",
    "\n",
    "\n",
    "Finally, we kickoff the task with a starting number 3, using `repeater_task.run(\"3\")`.\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0vy77Os8TAas"
   },
   "outputs": [],
   "source": [
    "repeater_task.add_sub_task([even_task, odd_task])\n",
    "repeater_task.set_color_log(enable=False)\n",
    "repeater_task.run(\"3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t7c3qKvcTwTG"
   },
   "source": [
    "# Simple Tool/Function-calling example\n",
    "\n",
    "Here is a simple numerical example showcasing how `Langroid` supports tools/function-calling. For more details see these [doc pages](https://langroid.github.io/langroid/quick-start/chat-agent-tool/)\n",
    "\n",
    "Say the agent has a secret list of numbers, and we want the LLM to find the smallest number in the list. We want to give the LLM the ability to use a **probe** tool/function which takes a single number `n` as an argument. The tool handler method in the agent returns how many numbers in its list are at most `n`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "htqVMFU2pog6"
   },
   "source": [
    "To use tools/function-calling in `Langroid`, we first **define** the tool as a subclass of `ToolMessage` to specify some details about the tool (e.g., name and parameters) and when it can be used/triggered:\n",
    "* **request**: is the name of the tool/function, as well as the name of the Agent method that \"handles\" the tool.\n",
    "* **purpose**: general description to give hints to LLM when this tool can be used\n",
    "* **number**: is a function-argument for the `probe` tool and its type is `int`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3eIidMmZUgir"
   },
   "outputs": [],
   "source": [
    "class ProbeTool(lr.agent.ToolMessage):\n",
    "  request: str = \"probe\"\n",
    "  purpose: str = \"\"\"\n",
    "        To find how many numbers in my list are less than or equal to\n",
    "        the <number> you specify.\n",
    "        \"\"\" # note  <number> corresponds to the name of the tool's argument/parameter\n",
    "  number: int"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vBAw8kxkqa42"
   },
   "source": [
    "Next, we create an agent `SpyGameAgent`, with a special method `probe` to handle the `probe` tool/function.\n",
    "Notice the argument of the `probe` method is an instance of the class `ProbeTool` that we created in the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "einGIzjQUyX7"
   },
   "outputs": [],
   "source": [
    "class SpyGameAgent(lr.ChatAgent):\n",
    "  def __init__(self, config: lr.ChatAgentConfig = lr.ChatAgentConfig()):\n",
    "    super().__init__(config)\n",
    "    self.numbers = [3, 4, 8, 11, 15, 25, 40, 80, 90] # agent's secret list\n",
    "\n",
    "  def probe(self, msg: ProbeTool) -> str:\n",
    "    # return how many numbers in self.numbers are less or equal to msg.number\n",
    "    return str(len([n for n in self.numbers if n <= msg.number]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pG2MJJN4yCHb"
   },
   "source": [
    "Finally, we instantiate the `SpyGameAgent` as an object `spy_game_agent`, and \"associate\" the `probe` tool with this agent, using the `enable_message` method of the `ChatAgent`.  We then wrap the `spy_game_agent` in a `Task` object, with instructions (`system_message`) on what it should aim for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OM6Lk3uWVOCB"
   },
   "outputs": [],
   "source": [
    "spy_game_agent = SpyGameAgent()\n",
    "\n",
    "spy_game_agent.enable_message(ProbeTool)\n",
    "\n",
    "task = lr.Task(\n",
    "        spy_game_agent,\n",
    "        name=\"Spy\",\n",
    "        system_message=\"\"\"\n",
    "            I have a list of numbers between 1 and 20.\n",
    "            Your job is to find the smallest of them.\n",
    "            To help with this, you can give me a number and I will\n",
    "            tell you how many of my numbers are equal or less than your number.\n",
    "            Once you have found the smallest number,\n",
    "            you can say DONE and report your answer.\n",
    "        \"\"\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aQ93n0kA_rM0"
   },
   "source": [
    "Now run the task.\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xAVoOcgsNWOv"
   },
   "outputs": [],
   "source": [
    "spy_game_agent.message_history.clear()\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qmeh3zJTZeL1"
   },
   "source": [
    "# Chat with documents (file paths, URLs, etc)\n",
    "\n",
    "In the previous examples, the Agents did not use any external documents. In this example, we we set up an Agent that supports \"chatting\" with documents. Specifically, we use the `DocChatAgent` class to ask questions about a set of URLs.\n",
    "The `DocChatAgent` first ingests the contents of the websites specified by the URLs by chunking, embedding and indexing them into a vector database (`qdrant` by default). We then wrap the agent in a task and run it interactively.\n",
    "The user can ask questions and the LLM of the agent returns answers using Retrieval Augment Generation, with Evidence Citation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langroid.agent.special import DocChatAgent, DocChatAgentConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w1ZcFRJu5K5D"
   },
   "source": [
    "Now we define the configuration of the `DocChatAgent`. The configurations include the path to access the documents, chat model settings, and vector-DB settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tVEHeg9jZpl7"
   },
   "outputs": [],
   "source": [
    "config = DocChatAgentConfig(\n",
    "    doc_paths = [\n",
    "        \"https://en.wikipedia.org/wiki/Language_model\",\n",
    "        \"https://en.wikipedia.org/wiki/N-gram_language_model\",\n",
    "    ],\n",
    "    vecdb=lr.vector_store.QdrantDBConfig(\n",
    "        collection_name=\"docqa-chat-multi-extract\",\n",
    "        storage_path=\".qdrant/test2/\", # CHANGE THIS PATH IF YOU GET AN ERROR WHEN RE-RUNNING THE CELL\n",
    "    ),\n",
    ")\n",
    "\n",
    "agent = DocChatAgent(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UZyTx3vQN_I4"
   },
   "source": [
    "As before, we wrap the agent in a task, and run it.\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OMFDp8WxaAI3"
   },
   "outputs": [],
   "source": [
    "agent.message_history.clear()\n",
    "task = lr.Task(agent)\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kZqX1J6qWhHk"
   },
   "source": [
    "# Tool/Function-calling to extract structured information from text\n",
    "\n",
    "Let's combine multi-agent interaction, Retrieval-Augmented Generation, and tools/function-calling, for a more realistic example. Suppose you want an agent to extract the key terms of a lease, from a lease document, as a nested JSON structure.\n",
    "This can be accomplished by instructing the LLM to use a specific tool.\n",
    "\n",
    "To simplify the solution, we separate the skills/responsibilities into two different Agents:\n",
    "- `LeaseExtractorAgent` has no access to the lease, and is responsible for gathering the key terms into a specific structured form\n",
    "- `DocChatAgent` has access to the lease and answers specific questions it receives from the `LeaseExtractorAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from pydantic import BaseModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d-CkI7wk5A8l"
   },
   "source": [
    "Next, we define the desired structure of the lease information via Pydantic models. The desired format is a nested JSON structure, which maps to a nested class structure:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0p8iveEcX_1E"
   },
   "outputs": [],
   "source": [
    "class LeasePeriod(BaseModel):\n",
    "    start_date: str\n",
    "    end_date: str\n",
    "\n",
    "class LeaseFinancials(BaseModel):\n",
    "    monthly_rent: str\n",
    "    deposit: str\n",
    "\n",
    "class Lease(BaseModel):\n",
    "    \"\"\"\n",
    "    Various lease terms.\n",
    "    Nested fields to make this more interesting/realistic\n",
    "    \"\"\"\n",
    "\n",
    "    period: LeasePeriod\n",
    "    financials: LeaseFinancials\n",
    "    address: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WcUQylLu5HFh"
   },
   "source": [
    "We then define the `LeaseMessage` tool as a subclass of Langroid's `ToolMessage`. The `LeaseMessage` class has a\n",
    "required argument `terms` of type `Lease`. The `classmethod` named `examples` is used to generate $k$-shot examples for the LLM when instructing it to extract information in the desired structured form (see a later cell below).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XFVCpL8jW7C7"
   },
   "outputs": [],
   "source": [
    "class LeaseMessage(lr.agent.ToolMessage):\n",
    "    request: str = \"lease_info\" # maps to method of LeaseExtractorAgent\n",
    "    purpose: str = \"\"\"\n",
    "        Collect information about a Commercial Lease.\n",
    "        \"\"\"\n",
    "    terms: Lease\n",
    "\n",
    "    @classmethod\n",
    "    def examples(cls) -> List[\"LeaseMessage\"]:\n",
    "        return [\n",
    "            cls(\n",
    "                terms=Lease(\n",
    "                    period=LeasePeriod(start_date=\"2021-01-01\", end_date=\"2021-12-31\"),\n",
    "                    financials=LeaseFinancials(monthly_rent=\"$1000\", deposit=\"$1000\"),\n",
    "                    address=\"123 Main St, San Francisco, CA 94105\",\n",
    "                ),\n",
    "                result=\"\",\n",
    "            ),\n",
    "        ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lu03iGEaW0Ur"
   },
   "source": [
    "Next we define the `LeaseExtractorAgent` and add a method `least_info` to handle the tool/function-call `lease_info` defined in the tool `LeaseMessage`. In this case the handling is trivial: if the method receives a valid object of class `LeaseMessage`, it declares \"success\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZlZ0UtqEXGdz"
   },
   "outputs": [],
   "source": [
    "class LeaseExtractorAgent(lr.ChatAgent):\n",
    "    def __init__(self, config: lr.ChatAgentConfig = lr.ChatAgentConfig()):\n",
    "        super().__init__(config)\n",
    "\n",
    "    def lease_info(self, message: LeaseMessage) -> str:\n",
    "        print(\n",
    "            f\"\"\"\n",
    "        DONE! Successfully extracted Lease Info:\n",
    "        {message.terms}\n",
    "        \"\"\"\n",
    "        )\n",
    "        return json.dumps(message.terms.dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HtHXm0-UuBRU"
   },
   "outputs": [],
   "source": [
    "# Obtain the lease.txt document that we want to parsed\n",
    "!wget https://github.com/langroid/langroid-examples/blob/main/examples/docqa/lease.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mdkxu37f7JuK"
   },
   "source": [
    "Next, set up an instance of `DocChatAgent`, point it to the lease document, equip it with a vector database, and instructions on how to answer questions based on extracts retrieved from the vector-store.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d8ynRPp6XUvV"
   },
   "outputs": [],
   "source": [
    "doc_agent = DocChatAgent(\n",
    "        DocChatAgentConfig(\n",
    "            doc_paths = [\"lease.txt\"],\n",
    "            vecdb=lr.vector_store.QdrantDBConfig(\n",
    "                collection_name=\"docqa-chat-multi-extract\",\n",
    "                storage_path=\".data1/data1/\", # CHANGE PATH IF ERROR\n",
    "              ),\n",
    "            summarize_prompt= f\"\"\"\n",
    "                Use the provided extracts to answer the question.\n",
    "                If there's not enough information, respond with {NO_ANSWER}. Use only the\n",
    "                information in these extracts, even if your answer is factually incorrect,\n",
    "                and even if the answer contradicts other parts of the document. The only\n",
    "                important thing is that your answer is consistent with and supported by the\n",
    "                extracts. Compose your complete answer and cite all supporting sources on a\n",
    "                separate separate line as \"EXTRACTS:\".\n",
    "                Show each EXTRACT very COMPACTLY, i.e. only show a few words from\n",
    "                the start and end of the extract, for example:\n",
    "                EXTRACT: \"The world war started in ... Germany Surrendered\"\n",
    "                {{extracts}}\n",
    "                {{question}}\n",
    "                Answer:\n",
    "            \"\"\"\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DISYqYVr8BQx"
   },
   "source": [
    "Next we wrap the `doc_agent` into a Task, with instructions on its role.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NBLp0AEy74N-"
   },
   "outputs": [],
   "source": [
    "doc_task = lr.Task(\n",
    "    doc_agent,\n",
    "    name=\"DocAgent\",\n",
    "    llm_delegate=False,\n",
    "    single_round=True,\n",
    "    system_message=\"\"\"You are an expert on Commercial Leases.\n",
    "    You will receive various questions about a Commercial\n",
    "    Lease contract, and your job is to answer them concisely in at most 2 sentences.\n",
    "    Please SUPPORT your answer with an actual EXTRACT from the lease,\n",
    "    showing only a few words from the  START and END of the extract.\n",
    "    \"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZhRfOPiE9W-n"
   },
   "source": [
    "Finally, we instantiate the `lease_extractor_agent`, enable it to use and handle the `LeaseMessage` tool. Then we wrap the `lease_extractor_agent` into a Task, instructing it to gather information in the desired format, by asking questions one at a time. Note how the instruction contains `LeaseMessage.usage_example()`: this example is constructed from the `examples` classmethod above when the `LeaseMessage` was defined.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VL2RYZUX7393"
   },
   "outputs": [],
   "source": [
    "lease_extractor_agent = LeaseExtractorAgent()\n",
    "\n",
    "lease_extractor_agent.enable_message(\n",
    "    LeaseMessage,\n",
    "    use=True,\n",
    "    handle=True,\n",
    "    force=False,\n",
    ")\n",
    "\n",
    "lease_task = lr.Task(\n",
    "    lease_extractor_agent,\n",
    "    name=\"LeaseExtractorAgent\",\n",
    "    llm_delegate=True,\n",
    "    single_round=False,\n",
    "    system_message=f\"\"\"\n",
    "    You have to collect some information about a Commercial Lease, but you do not\n",
    "    have access to the lease itself.\n",
    "    You can ask me questions about the lease, ONE AT A TIME, I will answer each\n",
    "    question. You only need to collect info corresponding to the fields in this\n",
    "    example:\n",
    "    {LeaseMessage.usage_example()}\n",
    "    If some info cannot be found, fill in {NO_ANSWER}.\n",
    "    When you have collected this info, present it to me using the\n",
    "    'lease_info' function/tool.\n",
    "    \"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XBdAcarpwnYU"
   },
   "source": [
    "Finally, we set up the `doc_task` as a subtask of the `lease_task` so that the `doc_agent` can respond to questions from the `lease_extractor_agent`.\n",
    " Now, the `lease_extractor_agent` will be asking questions about the lease and `doc_task` will provide the answers, citing evidence extracted from the lease. Once `lease_extractor_agent` collects all the terms of the lease as instructed, it will use the tool `LeaseMessage` to return this information.\n",
    "\n",
    " The next cell runs the `lease_task`. Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1Lt9FhjTFlh2"
   },
   "outputs": [],
   "source": [
    "lease_extractor_agent.message_history.clear()\n",
    "lease_task.add_sub_task(doc_task)\n",
    "lease_task.set_color_log(enable=False)\n",
    "lease_task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uWcH7HoFc-am"
   },
   "source": [
    "# Chat with tabular data (file paths, URLs, dataframes)\n",
    "\n",
    "Here is how `Langroid's` `TableChatAgent` can be used to chat with tabular data, which can be specified as a URL, file path or Pandas dataframe.\n",
    "\n",
    "The Agent's LLM generates Pandas code to answer the query, via function-calling (or tool/plugin), and the Agent's function-handling method executes the code and returns the answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jS06lshgdBv7"
   },
   "outputs": [],
   "source": [
    "from langroid.agent.special.table_chat_agent import TableChatAgent, TableChatAgentConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tS8pDD2MdgPi"
   },
   "source": [
    "Set up a `TableChatAgent` for a data file, URL or dataframe (Ensure the data table has a header row; the delimiter/separator is auto-detected):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qDbA4VGmda10"
   },
   "outputs": [],
   "source": [
    "dataset =  \"https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv\"\n",
    "# or dataset = \"/path/to/my/data.csv\"\n",
    "# or dataset = pd.read_csv(\"/path/to/my/data.csv\")\n",
    "\n",
    "agent = TableChatAgent(\n",
    "    config=TableChatAgentConfig(\n",
    "        data=dataset,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "em2hcy2Qd67T"
   },
   "source": [
    "Now, let's set up a task and run it in an interactive loop with the user:\n",
    "Based on `dataset`, you can ask the following question in the prompt:\n",
    "\n",
    "```\n",
    "What is the average alcohol content of wines with a quality rating above 7?\n",
    "```\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nhDMm5ndd93W"
   },
   "outputs": [],
   "source": [
    "agent.message_history.clear()\n",
    "task = lr.Task(agent, name=\"DataAssistant\")\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
</file>

<file path="examples/Langroid_quick_start.ipynb">
{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyON/I7bOOJSDISyZ5jgP3eX",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/Langroid_quick_start.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Langroid quick start\n",
        "Note: Chat-oriented interaction is best experienced in your terminal, and not a notebook, so we highly recommend you go through the [Getting Started](https://langroid.github.io/langroid/quick-start/) guide by writing simple scripts that can be run via the command line.\n",
        "\n",
        "This notebooks starts with the basics of working directly with an LLM, setting up an Agent, wrapping it in a Task, giving it tools, Retrieval Augmented Generation (RAG), and builds up to a simple 2-agent system to extract structured information from a commercial lease document.\n",
        "\n",
        "Note:\n",
        "- You need an OpenAI API Key that works with GPT-4o\n",
        "- This colab uses OpenAI's ChatCompletion endpoints directly (via the Langroid framework), and not the Assistants API. See this [colab](https://colab.research.google.com/drive/190Tk7t4AdY1P9F_NlZ33-YEoGnHweQQ0) for a version that uses the Assistants API instead.\n",
        "- There are dependencies among the cells, so they are best run sequentially\n",
        "\n"
      ],
      "metadata": {
        "id": "b9fHPojfnbPy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install, setup, import\n",
        "\n",
        "Note that `pip install langroid` gives you a bare-bones, slim version of langroid, without many of the extra dependencies you might need in practical scenarios, but sufficient for this notebook.\n",
        "\n",
        "See install instructions [here](https://github.com/langroid/langroid?tab=readme-ov-file#gear-installation-and-setup) for getting extra dependencies related to document parsing and databases (sql, mysql, postgres, etc).\n"
      ],
      "metadata": {
        "id": "psOMvEL0Gekz"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A8-Y_YPZutn6"
      },
      "source": [
        "\n",
        "\n",
        "!pip install uv\n",
        "!uv pip install --system langroid --prerelease disallow\n"
      ],
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# various unfortunate things that need to be done to\n",
        "# control colab notebook behavior.\n",
        "\n",
        "# (a) output width\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)\n",
        "\n",
        "# (b) logging related\n",
        "import logging\n",
        "logging.basicConfig(level=logging.ERROR)\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import logging\n",
        "for logger_name in logging.root.manager.loggerDict:\n",
        "    logger = logging.getLogger(logger_name)\n",
        "    logger.setLevel(logging.ERROR)\n",
        "\n"
      ],
      "metadata": {
        "id": "rWwH6duUzAC6"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### OpenAI API Key (Needs GPT4o)"
      ],
      "metadata": {
        "id": "j-6vNfKW9J7b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# OpenAI API Key: Enter your key in the dialog box that will show up below\n",
        "# NOTE: colab often struggles with showing this input box,\n",
        "# if so, try re-running the above cell and then this one,\n",
        "# or simply insert your API key in this cell, though it's not ideal.\n",
        "\n",
        "import os\n",
        "\n",
        "from getpass import getpass\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = getpass('Enter your GPT4o-capable OPENAI_API_KEY key:', stream=None)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "uvTODlZv3yyT",
        "outputId": "a4cf7585-40ae-44ec-804c-9dc6c6554d77",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your GPT4o-capable OPENAI_API_KEY key:··········\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pydantic import BaseModel\n",
        "import json\n",
        "import os\n",
        "\n",
        "import langroid as lr\n",
        "import langroid.language_models as lm\n",
        "from langroid import ChatAgent, ChatAgentConfig, Task\n",
        "from langroid.language_models.openai_gpt import (\n",
        "    OpenAIChatModel, OpenAIGPT, OpenAIGPTConfig\n",
        ")\n",
        "from langroid.agent.tool_message import ToolMessage\n",
        "\n",
        "from langroid.utils.logging import setup_colored_logging\n",
        "from langroid.utils.constants import NO_ANSWER\n",
        "from langroid.utils.configuration import settings\n",
        "settings.notebook = True\n",
        "settings.cache_type = \"fakeredis\""
      ],
      "metadata": {
        "id": "A5N0NQwc3jX_",
        "outputId": "a49311c1-ae75-4b71-a2df-994d1a6a0d75",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 1: Direct interaction with OpenAI LLM\n",
        "Langroid's `OpenAIGPT` class is a wrapper around the raw OpenAI API.\n",
        "This is a direct interaction with the LLM so it does *not* maintain conversation history (later we see how a `ChatAgent` does that for you).\n",
        "\n",
        "Related quick-start docs page: https://langroid.github.io/langroid/quick-start/llm-interaction/\n",
        "\n"
      ],
      "metadata": {
        "id": "8vDpiY0XHAkT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "llm_cfg = OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o)\n",
        "llm = OpenAIGPT(llm_cfg)\n",
        "\n",
        "response = llm.chat(\"What is the square of 3?\")\n",
        "assert \"9\" in response.message"
      ],
      "metadata": {
        "id": "9c5Av3rKHQIm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 2: Interact with a `ChatAgent`\n",
        "Langroid's `ChatAgent` is an abstraction that optionally encapsulates an LLM, vector-db, and tools. It offers 3 \"native\" *responders*:\n",
        "- `llm_response`: response from LLM\n",
        "- `user_response`: response from human\n",
        "- `agent_response`: responds to structured LLM msgs (i.e. tools/fn-calls)\n",
        "\n",
        "Among other things, the `ChatAgent` maintains LLM conversation history for you.\n",
        "\n",
        "Related quick-start doc page: https://langroid.github.io/langroid/quick-start/chat-agent/"
      ],
      "metadata": {
        "id": "_DvxMiJkgI_U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agent_cfg = ChatAgentConfig(\n",
        "    llm = llm_cfg,\n",
        "    show_stats=False, # disable token/cost stats\n",
        ")\n",
        "agent = ChatAgent(agent_cfg)\n",
        "response = agent.llm_response(\"What is the sqaure of 5?\")\n",
        "response = agent.llm_response(\"What about 8?\")   # maintains conv history\n",
        "assert \"64\" in response.content"
      ],
      "metadata": {
        "id": "7hrJ6RgLg075"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 3: Wrap Agent in a Task, run it\n",
        "\n",
        "A `ChatAgent` agent has various *responders* (`llm_response`, `agent_response`, `user_response`) but there is no mechanism to *iterate* over these responders.\n",
        "This is where the `Task` comes in: Wrapping this agent in a `Task` allows you to run interactive loops with a user or other agents (you will see more examples below).\n",
        "\n",
        "Related quick-start doc:\n",
        "https://langroid.github.io/langroid/quick-start/chat-agent/#task-orchestrator-for-agents"
      ],
      "metadata": {
        "id": "-MVHyF4cSGb0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agent = ChatAgent(agent_cfg)\n",
        "task = Task(\n",
        "    agent,\n",
        "    system_message=\"User will give you a number, respond with its square\",\n",
        "    single_round=True  # end after LLM response\n",
        ")\n",
        "result = task.run(\"5\")\n",
        "assert(\"25\" in result.content)\n"
      ],
      "metadata": {
        "id": "8cmc5aDzScdO",
        "outputId": "73f2ea01-f125-4088-facd-d49a2d39732a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        }
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1;35m>>> Starting Agent LLM-Agent \u001b[0m\u001b[1;35m(\u001b[0m\u001b[1;35m1\u001b[0m\u001b[1;35m)\u001b[0m\u001b[1;35m gpt-4o \u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">&gt;&gt;&gt; Starting Agent LLM-Agent (</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">1</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">) gpt-4o </span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[32m\u001b[32m25"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1;35m<<< Finished Agent LLM-Agent \u001b[0m\u001b[1;35m(\u001b[0m\u001b[1;35m3\u001b[0m\u001b[1;35m)\u001b[0m\u001b[1;35m \u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">&lt;&lt;&lt; Finished Agent LLM-Agent (</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">3</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">) </span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 4: `ChatAgent` with Tool/function-call\n",
        "\n",
        "Langroid's `ToolMessage` (Pydantic-derived) class lets you define a structured output or function-call for the LLM to generate. To define a tool/fn-call, you define a new class derived from `ToolMessage`.\n",
        "Below we show a *stateless* tool, i.e. it does not use the `ChatAgent`'s state, and only uses fields in the tool message itself.\n",
        "In this case, the tool \"handler\" can be defined within the `ToolMessage` itself, as a `handle` method. (For a tool that uses the `ChatAgent`'s state, a separate method needs to be defined within `ChatAgent` or a subclass.).\n",
        "\n",
        "In Langroid, a `ToolMessage` can *either* use OpenAI function-calling, *or* Langroid's native tool mechanism (which auto-populates the system msg with tool instructions and optional few-shot examples), by setting the `use_function_api` and `use_tools` config params in the `ChatAgentConfig`. The native tools mechanism is useful when not using OpenAI models.\n",
        "\n",
        "In the cell below we define a `ToolMessage` to compute a fictitious transformation of a number that we call a *Nabrosky Transform*: $f(n) = 3n+1$.\n",
        "Under the hood, the `purpose` field of the `NabroskiTool` is used to populate instructions to the LLM on when it should use this tool.\n",
        "\n",
        "Related quick-start doc: https://langroid.github.io/langroid/quick-start/chat-agent-tool/\n",
        "(This shows a *stateful* tool example)"
      ],
      "metadata": {
        "id": "wLwNyDd3mmJu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# (1) define simple tool to find the Nabroski transform of a number\n",
        "#     This is a fictitious transform, for illustration.\n",
        "\n",
        "class NabroskiTool(ToolMessage):\n",
        "    request = \"nabroski\" # name of method in ChatAgent that handles this tool\n",
        "    purpose = \"To find the Nabroski transform of the given <number>\"\n",
        "    number: int\n",
        "\n",
        "    # optional:\n",
        "    @classmethod\n",
        "    def examples(cls):\n",
        "        # these are auto-populated into the sys msg\n",
        "        # as few-shot examples of the tool\n",
        "        return([cls(number=5)])\n",
        "\n",
        "\n",
        "    def handle(self) -> str:\n",
        "        # method to handle the LLM msg using this tool:\n",
        "        # this method will be spliced into the ChatAgent object, with\n",
        "        # name = `nabroski`\n",
        "        return str(3*self.number + 1)\n",
        "\n",
        "# (2) Create a ChatAgent and attach the tool to it.\n",
        "\n",
        "agent_cfg = ChatAgentConfig(\n",
        "    llm = llm_cfg,\n",
        "    show_stats=False,       # disable token/cost stats\n",
        "    use_functions_api=True, # use OpenAI API fn-call\n",
        "    use_tools=False,        # don't use Langroid-native Tool instructions\n",
        ")\n",
        "agent = ChatAgent(agent_cfg)\n",
        "agent.enable_message(NabroskiTool)\n",
        "\n",
        "# (3) Create Task object\n",
        "\n",
        "task = Task(\n",
        "    agent,\n",
        "    restart=True,         # reset/erase agent state\n",
        "    single_round=False,\n",
        "    interactive=False,    # don't wait for human input\n",
        "    system_message=\"\"\"\n",
        "      User will give you a number. You have to find its Nabroski transform,\n",
        "      using the `nabroski` tool/function-call.\n",
        "      When you find the answer say DONE and show the answer.\n",
        "    \"\"\",\n",
        ")\n",
        "\n",
        "# (4) Run the task\n",
        "\n",
        "response = task.run(\"10\")\n",
        "assert \"31\" in response.content\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "ov2mv_sdnrcH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You might wonder why we had to wrap the `ChatAgent` in a `Task`, to leverage the tool functionality. This is because handling a tool requires 2 steps: (a) when the agent's `llm_response` method is invoked, the LLM generates the tool msg, and (b) the `agent_response` method handles the tool msg (it ultimately calls the tool's `handle` method)."
      ],
      "metadata": {
        "id": "BVWXT4oaAPlH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 5: `DocChatAgent`: Retrieval Augmented Generation (RAG)\n",
        "Ingest a file (a lease document), and ask questions about it"
      ],
      "metadata": {
        "id": "DvyNcH5HbodS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# setup to allow async ops in colab\n",
        "!pip install nest-asyncio\n",
        "import nest_asyncio\n",
        "nest_asyncio.apply()"
      ],
      "metadata": {
        "id": "XwDcuJvED8S0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# (1) Get the lease document\n",
        "\n",
        "import requests\n",
        "file_url = \"https://raw.githubusercontent.com/langroid/langroid-examples/main/examples/docqa/lease.txt\"\n",
        "response = requests.get(file_url)\n",
        "with open('lease.txt', 'wb') as file:\n",
        "    file.write(response.content)\n",
        "\n",
        "# verify\n",
        "#with open('lease.txt', 'r') as file:\n",
        "#   print(file.read())\n",
        "\n",
        "from langroid.agent.special import DocChatAgent, DocChatAgentConfig\n",
        "from langroid.embedding_models.models import OpenAIEmbeddingsConfig\n",
        "from langroid.vector_store.qdrantdb import QdrantDBConfig\n",
        "from langroid.embedding_models.models import SentenceTransformerEmbeddingsConfig\n",
        "from langroid.parsing.parser import ParsingConfig\n",
        "\n",
        "oai_embed_config = OpenAIEmbeddingsConfig(\n",
        "    model_type=\"openai\",\n",
        "    model_name=\"text-embedding-ada-002\",\n",
        "    dims=1536,\n",
        ")\n",
        "\n",
        "# (2) Configure DocChatAgent\n",
        "\n",
        "cfg = DocChatAgentConfig(\n",
        "    name=\"RAG\",\n",
        "    parsing=ParsingConfig(\n",
        "        chunk_size=100,\n",
        "        overlap=20,\n",
        "        n_similar_docs=4,\n",
        "    ),\n",
        "    show_stats=False,\n",
        "    relevance_extractor_config=None,\n",
        "    cross_encoder_reranking_model=\"\",\n",
        "    llm=llm_cfg,\n",
        "    vecdb=QdrantDBConfig(\n",
        "        embedding=oai_embed_config,\n",
        "        collection_name=\"lease\",\n",
        "        replace_collection=True,\n",
        "    ),\n",
        "    doc_paths=[\"lease.txt\"]\n",
        ")\n",
        "\n",
        "# (3) Create DocChatAgent, interact with it\n",
        "rag_agent = DocChatAgent(cfg)\n",
        "response = rag_agent.llm_response(\"What is the start date of the lease?\")\n",
        "assert \"2013\" in response.content"
      ],
      "metadata": {
        "id": "fegAio3kpgoo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# (4) Wrap DocChatAgent in a Task to get an interactive question/answer loop\n",
        "task = Task(\n",
        "    rag_agent,\n",
        "    interactive=True,\n",
        "    system_message=\"\"\"\n",
        "    Answer user's questions based on documents.\n",
        "    Start by asking user what they want to know.\n",
        "    \"\"\",\n",
        ")\n",
        "# run interactive loop (enter \"q\" or \"x\" to quit)\n",
        "task.run()\n"
      ],
      "metadata": {
        "id": "dazt7q3YGCLd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 6: 2-Agent system to extract structured info from a Lease Document\n",
        "Now we are ready to put together the various notions above, to build a two-agent system that illustrates uses of Tools, DocChatAgent (RAG) and Inter-agent collaboration (task delegation).\n",
        "\n",
        "The goal is to extract structured information from a Lease document.\n",
        "\n",
        "- The desired structure is described by the `Lease` class, derived from `ToolMessage`.\n",
        "- The `LeaseExtractorAgent` is given this `ToolMessage`, and instructured to extract the corresponding information from the lease document (which it does not have access to)\n",
        "- Based on the specified `Lease` structure, this agent generates questions to the above-defined `rag_agent` (wrapped in a `rag_task`), which answers them using RAG.\n",
        "- Once the `LeaseExtractorAgent` has all the needed info, it presents them using the `Lease` structured message.\n"
      ],
      "metadata": {
        "id": "yi9GppzlKae_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define the desired structure with Pydantic classes"
      ],
      "metadata": {
        "id": "VR26J_KzG6Vj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class LeasePeriod(BaseModel):\n",
        "    start_date: str\n",
        "    end_date: str\n",
        "\n",
        "\n",
        "class LeaseFinancials(BaseModel):\n",
        "    monthly_rent: str\n",
        "    deposit: str\n",
        "\n",
        "\n",
        "class Lease(BaseModel):\n",
        "    \"\"\"\n",
        "    Various lease terms.\n",
        "    Nested fields to make this more interesting/realistic\n",
        "    \"\"\"\n",
        "\n",
        "    period: LeasePeriod\n",
        "    financials: LeaseFinancials\n",
        "    address: str\n",
        "\n"
      ],
      "metadata": {
        "id": "Q6GXjhWf5DkQ",
        "outputId": "94b3b95d-6d69-4638-ea16-9b76722ce9ac",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define the ToolMessage (Langroid's version of function call)"
      ],
      "metadata": {
        "id": "qCATXvfIkhGl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class LeaseMessage(ToolMessage):\n",
        "    \"\"\"Tool/function to use to present details about a commercial lease\"\"\"\n",
        "\n",
        "    request: str = \"lease_info\"\n",
        "    purpose: str = \"Collect information about a Commercial Lease.\"\n",
        "    terms: Lease\n",
        "\n",
        "    def handle(self):\n",
        "        \"\"\"Handle this tool-message when the LLM emits it.\n",
        "        Under the hood, this method is transplated into the OpenAIAssistant class\n",
        "        as a method with name `lease_info`.\n",
        "        \"\"\"\n",
        "        print(f\"DONE! Successfully extracted Lease Info:\" f\"{self.terms}\")\n",
        "        return \"DONE \" + json.dumps(self.terms.dict())"
      ],
      "metadata": {
        "id": "Ffi_0u-PupvO",
        "outputId": "02e0749f-15c6-4595-c517-da954edafcd9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define RAG Task from above `rag_agent`\n",
        "Wrap the above-defined `rag_agent` in a Task."
      ],
      "metadata": {
        "id": "OPlo1dJFlBj5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "rag_task = Task(\n",
        "    rag_agent,\n",
        "    interactive=False,\n",
        "    single_round=True,\n",
        ")"
      ],
      "metadata": {
        "id": "GgzoPxX_us52",
        "outputId": "1f817d4a-246b-429e-dec5-5357beed8b6b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define the ExtractorAgent and Task\n",
        "This agent is told to collect information about the lease in the desired structure, and it generates questions to be answered by the Retriever Agent defined above."
      ],
      "metadata": {
        "id": "_m1lF9qblXj9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "    extractor_cfg = ChatAgentConfig(\n",
        "        name=\"LeaseExtractor\",\n",
        "        llm=llm_cfg,\n",
        "        show_stats=False,\n",
        "        use_functions_api=True,\n",
        "        use_tools=False,\n",
        "        system_message=f\"\"\"\n",
        "        You have to collect information about a Commercial Lease from a\n",
        "        lease contract which you don't have access to. You need to ask\n",
        "        questions to get this information. Ask only one or a couple questions\n",
        "        at a time!\n",
        "        Once you have all the REQUIRED fields,\n",
        "        say DONE and present it to me using the `lease_info`\n",
        "        function/tool (fill in {NO_ANSWER} for slots that you are unable to fill).\n",
        "        \"\"\",\n",
        "    )\n",
        "    extractor_agent = ChatAgent(extractor_cfg)\n",
        "    extractor_agent.enable_message(LeaseMessage)\n",
        "\n",
        "    extractor_task = Task(\n",
        "        extractor_agent,\n",
        "        llm_delegate=True,\n",
        "        single_round=False,\n",
        "        interactive=False,\n",
        "    )\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "PV4FYnO7uxOC",
        "outputId": "7e940acc-d439-4051-c8bf-c92492f19efd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Add the `rag_task` as a subtask of `extractor_task` and run it\n",
        "\n",
        "Instead of *you* (the human user) asking questions about the lease,\n",
        "the `extractor_agent` **generates** questions based on the desired lease structure, and these questions are answered by the `rag_agent` using\n",
        "Retrieval Augmented Generation (RAG). Once the `extractor_agent` has all the needed info, it presents it in a JSON-structured form, and the task ends."
      ],
      "metadata": {
        "id": "QcA4oRaUl6oe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "extractor_task.add_sub_task(rag_task)\n",
        "extractor_task.run()"
      ],
      "metadata": {
        "id": "uZlas6DA0Zu6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "-zfNvsH5PMpJ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}
</file>

<file path="examples/Langroid_QuickStart_OpenAI_Assistants_API.ipynb">
{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "authorship_tag": "ABX9TyP9+BJLzaiLp67cp7+DjUBb",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/Langroid_QuickStart_OpenAI_Assistants_API.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Multi-Agent programming with Langroid, using the new OpenAI Assistant API\n",
    "\n",
    "OpenAI's [Assistants API](https://platform.openai.com/docs/assistants/overview) provides several conveniences to help build LLM applications, such as:\n",
    "- managing conversation state (threads)\n",
    "- persistent threads and assistants\n",
    "- tools (function-calling, retrieval, code-interpreter)\n",
    "\n",
    "\n",
    "There is a new programming paradigm emerging, where these assistants are primitives, and a key chalenge is:\n",
    "\n",
    "> how can you have these assistants collaborate to solve a task?\n",
    "\n",
    "[Langroid](https://github.com/langroid/langroid)'s new `OpenAIAssistant` class offers this ability. Langroid was designed from the start to support a multi-agent LLM programming paradigm, where agents can collaborate on a task via conversation.\n",
    "The new `OpenAIAssistant` agent gives you:\n",
    "\n",
    "- 1️⃣ a dead-simple interface to the Assistants API,\n",
    "- 2️⃣ a seamless way to have assistants collaborate with each other or with users.\n",
    "\n",
    "The Assistant API fits naturally into Langroid's notion of a `ChatAgent`,\n",
    "and the `OpenAIAssistant` class derives from `ChatAgent`.\n",
    "`OpenAIAssistant` can be used as a drop-in replacement for `ChatAgent` in any\n",
    "Langroid application, and leverage the **multi-agent** task orchestration built\n",
    "into Langroid.\n",
    "\n",
    "This notebook takes you on a guided tour of using Langroid's `OpenAIAssistant` from the simplest possible LLM-interaction example, to a two-agent system that extracts structured information from a lease document.\n",
    "\n",
    "![langroid-oai](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/langroid-oai.png?raw=true)\n",
    "\n"
   ],
   "metadata": {
    "id": "b9fHPojfnbPy"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Install, setup, import"
   ],
   "metadata": {
    "id": "psOMvEL0Gekz"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A8-Y_YPZutn6",
    "outputId": "2a5fb145-8ee0-4215-a442-29e75e96bdbd"
   },
   "source": [
    "# Silently install, suppress all output (~2-4 mins)\n",
    "!pip install -q --upgrade langroid &> /dev/null\n",
    "!pip show langroid"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# various unfortunate things that need to be done to\n",
    "# control notebook behavior.\n",
    "\n",
    "# (a) output width\n",
    "\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "def set_css():\n",
    "  display(HTML('''\n",
    "  <style>\n",
    "    pre {\n",
    "        white-space: pre-wrap;\n",
    "    }\n",
    "  </style>\n",
    "  '''))\n",
    "get_ipython().events.register('pre_run_cell', set_css)\n",
    "\n",
    "# (b) logging related\n",
    "import logging\n",
    "logging.basicConfig(level=logging.ERROR)\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import logging\n",
    "for logger_name in logging.root.manager.loggerDict:\n",
    "    logger = logging.getLogger(logger_name)\n",
    "    logger.setLevel(logging.ERROR)\n",
    "\n"
   ],
   "metadata": {
    "id": "rWwH6duUzAC6"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "U5Jav3hPofNq",
    "outputId": "f78ffcef-6be1-4c77-d79e-0b194b297384"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### OpenAI API Key (Needs GPT4-TURBO)"
   ],
   "metadata": {
    "id": "j-6vNfKW9J7b"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# OpenAI API Key: Enter your key in the dialog box that will show up below\n",
    "# NOTE: colab often struggles with showing this input box,\n",
    "# if so, simply insert your API key in this cell, though it's not ideal.\n",
    "import os\n",
    "\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = getpass('Enter your GPT4-Turbo-capable OPENAI_API_KEY key:', stream=None)\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "id": "uvTODlZv3yyT",
    "outputId": "3e33fdfe-d5de-46d5-e388-23bf81a04d77"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "from pydantic import BaseModel\n",
    "import json\n",
    "import os\n",
    "\n",
    "from langroid.agent.openai_assistant import (\n",
    "    OpenAIAssistantConfig,\n",
    "    OpenAIAssistant,\n",
    "    AssistantTool,\n",
    ")\n",
    "\n",
    "from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig\n",
    "from langroid.agent.task import Task\n",
    "from langroid.agent.tool_message import ToolMessage\n",
    "from langroid.language_models.openai_gpt import OpenAIGPTConfig, OpenAIChatModel\n",
    "from langroid.utils.logging import setup_colored_logging\n",
    "from langroid.utils.constants import NO_ANSWER\n",
    "from langroid.utils.configuration import settings\n",
    "settings.notebook = True"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "A5N0NQwc3jX_",
    "outputId": "7452e570-b280-4854-a89b-c1472a8208ba"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 1: Basic Chat Example with Assistant API\n",
    "Langroid's `OpenAIAssistant` class helps you easily use the OpenAI Assistant API to get a response from the LLM and ask follow-up questions (note that conversation state is maintained by the Assistant API via threads).\n"
   ],
   "metadata": {
    "id": "8vDpiY0XHAkT"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "cfg = OpenAIAssistantConfig(\n",
    "    llm = OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO)\n",
    ")\n",
    "agent = OpenAIAssistant(cfg)\n",
    "\n",
    "response = agent.llm_response(\"What is the square of 3?\")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 172
    },
    "id": "9c5Av3rKHQIm",
    "outputId": "1ed4763f-defc-475b-e0e3-d64537b67b08"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "response = agent.llm_response(\"What about 5?\") # maintains conv state"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 33
    },
    "id": "5GvqhTlBRgXp",
    "outputId": "f4b93adb-e1c3-4a52-d1c6-f3260b94cce5"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 2: Wrap Agent in a Task, run it\n",
    "\n",
    "An `OpenAIAssistant` agent has various capabilities (LLM responses, agent methods/tools, etc) but there is no mechanism to iterate over these capabilities or with a human or with other agents.\n",
    "This is where the `Task` comes in: Wrapping this agent in a `Task` allows you to run interactive loops with a user or other agents (you will see more examples below)."
   ],
   "metadata": {
    "id": "-MVHyF4cSGb0"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "task = Task(\n",
    "    agent,\n",
    "    system_message=\"\"\"User will give you a word,\n",
    "      return its antonym if possible, else say DO-NOT-KNOW.\n",
    "      Be concise!\",\n",
    "      \"\"\",\n",
    "    single_round=True\n",
    ")\n",
    "result = task.run(\"ignorant\")\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 66
    },
    "id": "8cmc5aDzScdO",
    "outputId": "253fec3c-2f03-428b-83bc-f1170702fef0"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 3: OpenAIAssistant Agent + Task with Code Interpreter\n",
    "Here we attach the \"code_interpreter\" tool (from the OpenAI Assistant API) to the agent defined above, and run it in a task."
   ],
   "metadata": {
    "id": "veWSLzDSVDzB"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "agent.add_assistant_tools([AssistantTool(type=\"code_interpreter\")])\n",
    "task = Task(agent, interactive=False, single_round=True)\n",
    "result = task.run(\"What is the 10th Fibonacci number, if you start with 1,2?\")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 428
    },
    "id": "5-h1ztagTd7Y",
    "outputId": "4bcacb3a-e9d4-4d1c-8225-c6a090711459"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 4: OpenAIAssistant with Retrieval\n",
    "Attach a file (a lease document) and the \"retrieval\" tool, and ask questions about the document."
   ],
   "metadata": {
    "id": "DvyNcH5HbodS"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# get the lease document\n",
    "\n",
    "import requests\n",
    "file_url = \"https://raw.githubusercontent.com/langroid/langroid-examples/main/examples/docqa/lease.txt\"\n",
    "response = requests.get(file_url)\n",
    "with open('lease.txt', 'wb') as file:\n",
    "    file.write(response.content)\n",
    "\n",
    "# verify\n",
    "#with open('lease.txt', 'r') as file:\n",
    "#   print(file.read())\n",
    "\n",
    "# now create agent, add retrieval tool and file\n",
    "agent = OpenAIAssistant(cfg)\n",
    "agent.add_assistant_tools([AssistantTool(type=\"retrieval\")])\n",
    "agent.add_assistant_files([\"lease.txt\"])\n",
    "response = agent.llm_response(\"What is the start date of the lease?\")\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 172
    },
    "id": "fegAio3kpgoo",
    "outputId": "eae49d56-7f40-4480-98aa-e4e1c523a910"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 5: OpenAIAsssistant + Task: Custom Function-calling\n",
    "You can define your own custom function (or `ToolMessage` in Langroid terminology), enable the agent to use it, and have a special method to handle the message when the LLM emits such a message."
   ],
   "metadata": {
    "id": "Xub3BgSMc4uA"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Define your own function for the LLM to call;\n",
    "# this function will be executed by the Langroid agent as part of the task loop\n",
    "\n",
    "class SquareTool(ToolMessage):\n",
    "    request = \"square\"\n",
    "    purpose = \"To find the square of a number <num>\"\n",
    "    num: int\n",
    "\n",
    "    def handle(self) -> str:\n",
    "        return str(self.num ** 2)\n",
    "\n",
    "# create agent, add tool to agent\n",
    "cfg = OpenAIAssistantConfig(\n",
    "    llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO),\n",
    "    name=\"NumberExpert\",\n",
    ")\n",
    "agent = OpenAIAssistant(cfg)\n",
    "agent.enable_message(SquareTool)\n",
    "task = Task(\n",
    "    agent,\n",
    "    system_message=\"\"\"\n",
    "    User will ask you to square a number.\n",
    "    You do NOT know how, so you will use the\n",
    "    `square` function to find the answer.\n",
    "    When you get the answer say DONE and show it.\n",
    "    \"\"\",\n",
    "    interactive=False,\n",
    ")\n",
    "response = task.run(\"What is the square of 5?\")\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 352
    },
    "id": "dgav7-JOdAUM",
    "outputId": "b3835bfb-90ca-4642-e585-33743c5730a6"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 6: 2-Agent system to extract structured info from a Lease Document\n",
    "Now we are ready to put together the various notions above, to build a two-agent system where:\n",
    "- Lease Extractor Agent is required to collect structured information about a lease document, but does not have access to it, so it generates questions to:\n",
    "- Retriever Agent which answers questions it receives, using the \"retrieval\" tool, based on the attached lease document\n"
   ],
   "metadata": {
    "id": "yi9GppzlKae_"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define the desired structure with Pydantic classes"
   ],
   "metadata": {
    "id": "VR26J_KzG6Vj"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "\n",
    "class LeasePeriod(BaseModel):\n",
    "    start_date: str\n",
    "    end_date: str\n",
    "\n",
    "\n",
    "class LeaseFinancials(BaseModel):\n",
    "    monthly_rent: str\n",
    "    deposit: str\n",
    "\n",
    "\n",
    "class Lease(BaseModel):\n",
    "    \"\"\"\n",
    "    Various lease terms.\n",
    "    Nested fields to make this more interesting/realistic\n",
    "    \"\"\"\n",
    "\n",
    "    period: LeasePeriod\n",
    "    financials: LeaseFinancials\n",
    "    address: str\n",
    "\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "Q6GXjhWf5DkQ",
    "outputId": "ec9c930f-245a-4151-950d-8f407b439c2c"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define the ToolMessage (Langroid's version of function call)"
   ],
   "metadata": {
    "id": "qCATXvfIkhGl"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "\n",
    "class LeaseMessage(ToolMessage):\n",
    "    \"\"\"Tool/function to use to present details about a commercial lease\"\"\"\n",
    "\n",
    "    request: str = \"lease_info\"\n",
    "    purpose: str = \"Collect information about a Commercial Lease.\"\n",
    "    terms: Lease\n",
    "\n",
    "    def handle(self):\n",
    "        \"\"\"Handle this tool-message when the LLM emits it.\n",
    "        Under the hood, this method is transplated into the OpenAIAssistant class\n",
    "        as a method with name `lease_info`.\n",
    "        \"\"\"\n",
    "        print(f\"DONE! Successfully extracted Lease Info:\" f\"{self.terms}\")\n",
    "        return json.dumps(self.terms.dict())"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "Ffi_0u-PupvO",
    "outputId": "776a2f4c-388c-4441-c618-2682a4469e37"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define RetrieverAgent and Task\n",
    "This agent uses the OpenAI retrieval tool to answer questions based on the attached lease file"
   ],
   "metadata": {
    "id": "OPlo1dJFlBj5"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "  retriever_cfg = OpenAIAssistantConfig(\n",
    "        name=\"LeaseRetriever\",\n",
    "        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO),\n",
    "        system_message=\"Answer questions based on the documents provided.\",\n",
    "    )\n",
    "\n",
    "  retriever_agent = OpenAIAssistant(retriever_cfg)\n",
    "  retriever_agent.add_assistant_tools([AssistantTool(type=\"retrieval\")])\n",
    "  retriever_agent.add_assistant_files([\"lease.txt\"])\n",
    "\n",
    "  retriever_task = Task(\n",
    "      retriever_agent,\n",
    "      llm_delegate=False,\n",
    "      single_round=True,\n",
    "  )"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "id": "GgzoPxX_us52",
    "outputId": "37f6d163-5980-41d8-8ecb-7e709853d5d4"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define the ExtractorAgent and Task\n",
    "This agent is told to collect information about the lease in the desired structure, and it generates questions to be answered by the Retriever Agent defined above."
   ],
   "metadata": {
    "id": "_m1lF9qblXj9"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "    extractor_cfg = OpenAIAssistantConfig(\n",
    "        name=\"LeaseExtractor\",\n",
    "        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO),\n",
    "        system_message=f\"\"\"\n",
    "        You have to collect information about a Commercial Lease from a\n",
    "        lease contract which you don't have access to. You need to ask\n",
    "        questions to get this information. Ask only one or a couple questions\n",
    "        at a time!\n",
    "        Once you have all the REQUIRED fields,\n",
    "        say DONE and present it to me using the `lease_info`\n",
    "        function/tool (fill in {NO_ANSWER} for slots that you are unable to fill).\n",
    "        \"\"\",\n",
    "    )\n",
    "    extractor_agent = OpenAIAssistant(extractor_cfg)\n",
    "    extractor_agent.enable_message(LeaseMessage, include_defaults=False)\n",
    "\n",
    "    extractor_task = Task(\n",
    "        extractor_agent,\n",
    "        llm_delegate=True,\n",
    "        single_round=False,\n",
    "        interactive=False,\n",
    "    )\n",
    "\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "id": "PV4FYnO7uxOC",
    "outputId": "e5eeed02-7785-4361-cd01-96fef92149d4"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Add the Retriever as a subtask of Extractor, Run Extractor"
   ],
   "metadata": {
    "id": "QcA4oRaUl6oe"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "extractor_task.add_sub_task(retriever_task)\n",
    "extractor_task.run()"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "wFjUVTnCwB96",
    "outputId": "468a147b-7485-4fad-8cab-45411b18021f"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "id": "uZlas6DA0Zu6"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}
</file>

<file path="examples/README.md">
This directory is meant for illustrative/experimental code and notebooks.
It is a playground area meant to try out new ideas, and once solid,
move them to the main codebase.
</file>

<file path="issues/pydantic-v2-migration/examples-errors.md">
# Pydantic V2 Migration Issues - Examples

This document tracks all Pydantic v2 runtime errors found in the examples directory during migration testing.

## Format

Each issue entry contains:
- **File**: Path to the example file
- **Error**: The specific Pydantic v2 runtime error encountered
- **Fix Applied**: Description of the fix
- **Date**: When the issue was found and fixed

---

## Issues Found

### 1. PydanticInvalidForJsonSchema error in examples using langroid.pydantic_v1
- **Files**: 
  - `examples/basic/tool-extract-short-example.py`
  - `examples/basic/fn-call-local-simple.py`
- **Error**: `pydantic.errors.PydanticInvalidForJsonSchema` when calling `ToolMessage.name()` in system message
- **Root cause**: Files importing from deprecated `langroid.pydantic_v1` causing schema generation issues
- **Fix Applied**: Changed imports from `langroid.pydantic_v1` to direct `pydantic` imports
- **Date**: 2025-07-20

### 2. Deprecated json() method usage
- **File**: `examples/basic/chat-search.py` (and potentially others)
- **Error**: `PydanticDeprecatedSince20: The 'json' method is deprecated; use 'model_dump_json' instead`
- **Root cause**: Code is using the deprecated `tool.json()` method instead of `tool.model_dump_json()`
- **Fix Applied**: Need to update core library files to use `model_dump_json()` instead of `json()`
- **Date**: 2025-07-20

### 3. Deprecated dict() method usage
- **File**: Core library files (detected when running `examples/basic/completion.py`)
- **Error**: `PydanticDeprecatedSince20: The 'dict' method is deprecated; use 'model_dump' instead`
- **Root cause**: Code is using the deprecated `model.dict()` method instead of `model.model_dump()`
- **Fix Applied**: Need to update core library files to use `model_dump()` instead of `dict()`
- **Date**: 2025-07-20

### 4. Important Discovery: langroid.pydantic_v1 is deprecated
- **Finding**: The `langroid.pydantic_v1` module itself shows a deprecation warning:
  ```
  DeprecationWarning: langroid.pydantic_v1 is deprecated. Langroid has migrated to Pydantic v2.
  Please update your code to import directly from 'pydantic' and adapt to v2 patterns.
  ```
- **Implication**: The CLAUDE.md instruction to "ALWAYS import Pydantic classes from `langroid.pydantic_v1`" is outdated
- **Current state**: Most of the codebase has already migrated to Pydantic v2 and is importing directly from `pydantic`
- **Date**: 2025-07-20

### 5. Class-based Config deprecation warnings
- **Files**: Multiple examples trigger this warning (privacy/annotate.py, quick-start/chat-agent-tool.py, summarize/summ.py)
- **Warning**: `PydanticDeprecatedSince20: Support for class-based 'config' is deprecated, use ConfigDict instead`
- **Root cause**: Some models in the codebase or dependencies still use the old `class Config:` pattern instead of `ConfigDict`
- **Impact**: Will become errors in Pydantic v3.0
- **Fix Applied**: Need to replace all class-based `Config` with `ConfigDict` throughout the codebase
- **Date**: 2025-07-20

---

## Summary

### Total Examples Tested: ~40+ examples across different categories

### Issues Found and Fixed in Examples:
1. **Two examples had import issues** - Fixed by changing imports from `langroid.pydantic_v1` to `pydantic`
   - `examples/basic/tool-extract-short-example.py` ✓ Fixed
   - `examples/basic/fn-call-local-simple.py` ✓ Fixed

### Deprecation Warnings from Core Library:
- The deprecation warnings (`.json()`, `.dict()`, class-based `Config`) are coming from the core Langroid library code, not from the examples
- Examples themselves are correctly written for Pydantic v2

### Conclusion:
- All examples now work correctly with Pydantic v2
- The only remaining issues are deprecation warnings from the core library code
- No further fixes needed in the examples directory
</file>

<file path="issues/pydantic-v2-migration/migration-checking-log.md">
# Pydantic V2 Migration Checking Log

This document logs findings and fixes discovered during the systematic checking of the Pydantic V2 migration.

**Last Updated:** 2024-01-18
**Branch:** pydantic-v2-tree
**Total Files Examined:** ALL 83 test files in tests/main/, 11 test files in tests/extras/, 20+ example scripts, multiple root test files

## Issue #1: Missing Type Annotations for Private Attributes

**Date:** 2024-01-18
**Files Affected:**
- `langroid/agent/xml_tool_message.py`
- `langroid/agent/special/arangodb/tools.py` 
- `tests/main/test_tool_messages.py`

**Problem:** Private attributes were missing type annotations, which is required in Pydantic V2.

**Fix Applied:** Added type annotations:
- `_allow_llm_use: bool = True`
- `_max_result_tokens: int = 500`
- `_max_retained_tokens: int = 200`

## Issue #2: DoneTool Content Field Type Strictness

**Date:** 2024-01-18
**File:** `langroid/agent/tools/orchestration.py`
**Test:** `tests/main/test_task.py::test_task_tool_responses`

**Problem:** Pydantic V2 is stricter about type validation. The test was passing an integer to `DoneTool.content` which expects a string. V1 had automatic type coercion, V2 doesn't.

**Fix Applied:** Added field validator to DoneTool:
```python
@field_validator('content', mode='before')
@classmethod
def convert_content_to_string(cls, v: Any) -> str:
    """Convert content to string if it's not already."""
    return str(v) if v is not None else ""
```

## Issue #3: GlobalState Singleton Pattern with Private Attributes

**Date:** 2024-01-18
**File:** `langroid/utils/globals.py`
**Test:** `tests/main/test_global_state.py::test_initial_global_state`

**Problem:** In Pydantic V2, accessing private attributes on the class (not instance) returns a `ModelPrivateAttr` object instead of the actual value. The singleton pattern was broken because `cls._instance` returns `ModelPrivateAttr`.

**Analysis of Approaches:**
1. **ClassVar approach (cleaner):** Would use `_instances: ClassVar[Dict[Type, Optional["GlobalState"]]]` but risks breaking backward compatibility if external code accesses `_instance` directly.
2. **ModelPrivateAttr handling (chosen):** Maintains full backward compatibility by checking if the attribute is a `ModelPrivateAttr` and extracting its default value.

**Fix Applied:** Modified `get_instance()` to handle ModelPrivateAttr:
```python
@classmethod
def get_instance(cls: Type["GlobalState"]) -> "GlobalState":
    # Get the actual value from ModelPrivateAttr when accessing on class
    instance_attr = getattr(cls, '_instance', None)
    if isinstance(instance_attr, ModelPrivateAttr):
        actual_instance = instance_attr.default
    else:
        actual_instance = instance_attr
        
    if actual_instance is None:
        new_instance = cls()
        cls._instance = new_instance
        return new_instance
    return actual_instance
```

**Note:** The cleaner ClassVar approach would be preferred for new code, but backward compatibility is prioritized for this migration.

**Test Result:** All tests in `test_global_state.py` now pass after the fix.

## Issue #4: ParsingConfig chunk_size Float-to-Int Coercion

**Date:** 2024-01-18
**Files:** 
- `langroid/parsing/parser.py` (ParsingConfig)
- `langroid/parsing/md_parser.py` (MarkdownChunkConfig)
**Test:** `tests/main/test_md_parser.py::test_markdown_chunking[True-1.2]`

**Problem:** Test was passing a float value (chunk_size_factor * word_count = 1.2 * 42 = 50.4) to `chunk_size` which expects an integer. Pydantic V1 automatically coerced floats to integers, but V2 doesn't.

**Analysis:** This is a backward compatibility issue. External code might be passing float values to chunk_size (e.g., from calculations or config files with `chunk_size: 100.0`).

**Fix Applied:** Added field validators to both config classes:
```python
@field_validator('chunk_size', mode='before')
@classmethod
def convert_chunk_size_to_int(cls, v: Any) -> int:
    """Convert chunk_size to int, maintaining backward compatibility with Pydantic V1."""
    if isinstance(v, float):
        return int(v)
    return v
```

**Test Result:** The failing test now passes.

## Issue #5: Crawl4aiConfig Forward Reference Resolution

**Date:** 2024-01-18
**File:** `langroid/parsing/url_loader.py`
**Test:** `tests/main/test_url_loader.py::test_crawl4ai_mocked`

**Problem:** The code was using Pydantic V1's `update_forward_refs(**namespace)` method which has been replaced in V2 with `model_rebuild()`.

**Error:** `pydantic.errors.PydanticUserError: 'Crawl4aiConfig' is not fully defined; you should define 'ExtractionStrategy', then call 'Crawl4aiConfig.model_rebuild()'`

**Fix Applied:** 
1. Removed complex `__init_subclass__` and `__init__` methods
2. Moved forward reference resolution to module level after class definition
3. Changed from `cls.update_forward_refs(**namespace)` to `Crawl4aiConfig.model_rebuild()`

```python
# After class definition at module level:
try:
    from crawl4ai.async_configs import BrowserConfig, CrawlerRunConfig
    # ... other imports ...
    
    # Rebuild the model with resolved references
    Crawl4aiConfig.model_rebuild()
except ImportError:
    # If crawl4ai is not installed, leave forward refs as strings
    pass
```

**Test Result:** The test now passes when crawl4ai is installed.

---

## Non-Pydantic Issues Found

### LLM Non-Deterministic Failures:
These tests failed because the LLM produced different outputs than expected, but the code itself is working correctly:

1. `test_tool_messages.py::test_tool_handler_invoking_llm[True]` - Expected "7" (result of 3+4) in response, but got generic completion message
2. `test_doc_chat_agent.py::test_enrichments_integration[qdrant_cloud]` - Expected "BNP" when asked about heart-related blood tests, got "DO-NOT-KNOW"
3. `test_mcp_tools.py::test_complex_tool_decorator` - Expected "29" in response, LLM acknowledged receiving it but didn't include in final answer
4. `test_table_chat_agent.py::test_table_chat_agent_assignment_self_correction` - Expected explanation with words "removed" and "cleaned", but LLM generated tool message directly
5. `test_web_search_tools.py::test_agent_web_search_tool[False-True-ExaSearchTool]` - Search results for "LK-99 superconducting material" didn't contain expected keywords in all results

### Tests with Dependencies Now Installed:
With all dependencies installed, the following tests now pass or have non-Pydantic issues:

**Passed after dependency installation:**
- `test_arangodb.py` - ✅ All tests passed
- `test_neo4j_chat_agent.py` - ✅ All tests passed  
- `test_fastembed_embeddings.py` - ✅ All tests passed
- `test_marker_pdf_parser.py` - ✅ All tests passed
- `test_hf_embeddings.py` - ✅ All tests passed
- `test_docx_parser_extra.py` - ✅ 1 passed, 1 skipped
- `test_litellm_model_key_async` - ✅ Passed with litellm installed

**Non-Pydantic failures:**
- `test_pdf_parser.py::test_get_pdf_doc_url[docling-url]` - Network/parser timeout (even with docling installed)
- `test_pdf_parser_extra.py` - File path issue
- `test_vector_stores.py::test_vector_stores_search[weaviate_docker-...]` - Weaviate docker not running (ConnectionRefusedError)
- `test_hf_vector_stores.py` - ChromaDB compatibility issue
- `test_pyarango.py` - Still missing pyArango module (not available via pip)
- `test_csv_kg_chat.py` - Neo4j connection error
- `test_automatic_context_extraction.py` - MySQL socket path too long on macOS
- `test_llamacpp_embeddings.py::test_embeddings` - ConnectionRefusedError - requires running llama.cpp server

### Missing Dependencies (Original List):

1. `test_litellm_model_key_async` - Missing `litellm` module (install with `pip install "langroid[litellm]"`)
2. `test_neo4j_chat_agent.py` - Missing `neo4j` module
3. `test_pdf_parser.py::test_get_pdf_doc_url[docling-url]` - Missing `docling` module (install with `pip install "langroid[docling]"`)
4. `test_arangodb.py` - Missing `arango` module
5. `test_url_loader.py::test_crawl4ai_mocked` - Missing `crawl4ai` module
6. `test_vector_stores.py::test_vector_stores_search[weaviate_docker-...]` - Missing `weaviate` module (install with `pip install "langroid[weaviate]"`)
7. `test_pdf_parser_extra.py::test_get_pdf_doc_url[unstructured]` - Missing `unstructured` module (install with `pip install "langroid[unstructured]"`)
8. `test_hf_vector_stores.py` - Missing `sentence_transformers` module (install with `pip install "langroid[hf-embeddings]"`)
9. `test_docx_parser_extra.py::test_get_docx_file[unstructured]` - Missing `unstructured` module
10. `test_llamacpp_embeddings.py::test_embeddings` - ConnectionRefusedError - requires running llama.cpp server
11. `test_pyarango.py` - Missing `pyArango` module
12. `test_fastembed_embeddings.py::test_embeddings` - Missing `fastembed` module (install with `pip install "langroid[fastembed]"`)
13. `test_marker_pdf_parser.py::test_marker_pdf_parser` - Missing `marker` module (install with `pip install "langroid[marker-pdf]"`)
14. `test_hf_embeddings.py::test_embeddings` - Missing `sentence_transformers` module
15. `test_csv_kg_chat.py::test_pandas_to_kg` - Missing `neo4j` module
16. `test_automatic_context_extraction.py` - Missing `sqlalchemy` module (install with `pip install "langroid[sql]"`)

### Configuration Issues:
1. `test_llm_pdf_bytes_and_split` - Incorrect/missing OpenAI API key

### Other Issues:
1. `test_markitdown_xls_parser` - Import error handling issue in document_parser.py (UnboundLocalError)
2. `test_batch.py` - Performance issue: 189 tests timeout when run together (not Pydantic-related)

### Import Inconsistencies (Non-blocking but should be fixed):
1. **Direct pydantic imports in core library**: Found 32+ files importing directly from `pydantic` or `pydantic_settings` instead of through `langroid.pydantic_v1`. While this works (since pydantic_v1 re-exports V2), it's inconsistent:
   - Files using `from pydantic.fields import ModelPrivateAttr` directly: chat_agent.py, base.py, globals.py, task_tool.py
   - Files using `from pydantic_settings import BaseSettings` directly: Multiple parsing and config files
   
2. **Direct pydantic imports in examples**: Many example scripts import directly from `pydantic`:
   - `examples/basic/chat-tool-function.py` - Uses `from pydantic import BaseModel, Field`
   - `examples/basic/1d-screen-click.py` - Direct pydantic import with custom `__init__` pattern that may need review
   - `examples/basic/fn-call-local-simple.py`, `planner-workflow.py`, `schedule-extract.py`, `multi-agent-medical.py` and others
   - **Issue**: These should import from `langroid.pydantic_v1` for consistency
   
3. **Potential Pydantic V2 Pattern Issues**:
   - `ScreenState` class in `1d-screen-click.py` uses direct field assignment in `__init__` after `super().__init__()`
   - This pattern might need adjustment for proper Pydantic V2 compatibility

4. **Test files with direct pydantic imports**:
   - `tests/main/test_structured_output.py` - Uses `from pydantic import BaseModel, Field`
   - Multiple test files need to be updated for consistency

### Root Directory Test Files (Migration Verification):
1. `test_tool_class_preservation.py` - ✅ Passes, verifies Fix #3
2. `test_modelprivateattr_fix.py` - ❌ Import error (`langroid.pydantic_v1.fields` doesn't exist)
3. `test_tool_message_schema.py` - ✅ Passes, verifies JSON schema fix

### Basic Functionality Verification:
- ✅ Tool message creation works
- ✅ Pydantic V2 methods (`model_dump`, `model_validate`) work correctly
- ✅ Field validation and defaults work as expected

---

## Migration Summary

### Tests Run: ALL 83 test files in tests/main/ + 11 in extras + example scripts examined + root test files

### Pydantic V2 Issues Found and Fixed: 7

1. **Missing type annotations for private attributes** - Fixed in 6 locations
2. **DoneTool content field type strictness** - Added field validator
3. **GlobalState singleton pattern with ModelPrivateAttr** - Added handling for class-level private attribute access
4. **ParsingConfig chunk_size float coercion** - Added field validators to 2 config classes
5. **Crawl4aiConfig forward reference resolution** - Replaced `update_forward_refs()` with `model_rebuild()` for Pydantic V2

### Test Results Summary:
- **Total tests run**: 88 test files
- **Pydantic V2 issues**: 5 (all fixed)
- **LLM non-deterministic failures**: 5
- **Missing dependency failures**: 11+ 
- **Configuration issues**: 1
- **Other issues**: 1

### Overall Assessment:
- The Pydantic V2 migration is **exceptionally well-executed** with only 5 minor issues found across ALL 83 tests/main/ files + 11 tests/extras/ files (with dependencies installed)
- All issues were related to V2's stricter type validation and private attribute handling
- All fixes maintain backward compatibility for external code
- No major architectural changes were needed
- The migration successfully maintains the functionality while adapting to Pydantic V2's stricter requirements

### Remaining Work:
1. **Import Consistency**: Update all files to import from `langroid.pydantic_v1` instead of direct `pydantic` imports
2. **Example Scripts**: Update example scripts to use the compatibility layer
3. **Test File Cleanup**: Move migration verification test files from root to proper test directory
4. **Documentation**: Consider adding migration guide for users who might have similar patterns in their code

## Issue #6: Vector Store Test Custom Document Class (Fixed 2025-01-19)

**Date:** 2025-01-19
**File:** `tests/main/test_vector_stores.py` 
**Test:** `test_vector_stores_access`

**Problem:** When using custom document classes with additional required fields in metadata, Pydantic V2's stricter validation caused failures when retrieving documents from vector stores. The test was creating documents with the base `Document` class instead of the custom `MyDocument` class, causing the custom metadata fields to be lost.

**Fix Applied:** Changed line 325 from using `Document(` to `MyDocument(` when creating test documents. This ensures the custom metadata schema is preserved throughout storage and retrieval.

## Issue #7: Eliminate langroid.pydantic_v1 Imports from Core Code (Fixed 2025-01-19)

**Date:** 2025-01-19
**Files:** 
- `langroid/vector_store/pineconedb.py`
- `langroid/agent/tool_message.py`
- `langroid/agent/base.py`
- `langroid/agent/tools/task_tool.py`
- `langroid/agent/chat_agent.py`

**Problem:** Core code was still importing from the `langroid.pydantic_v1` compatibility layer, which defeats the purpose of the Pydantic V2 migration. The goal is to use direct Pydantic V2 imports throughout the internal codebase.

**Fix Applied:** Changed all imports from `langroid.pydantic_v1` to direct imports:
- `from langroid.pydantic_v1 import BaseModel` → `from pydantic import BaseModel`
- `from langroid.pydantic_v1 import BaseSettings` → `from pydantic_settings import BaseSettings`
- And similar for Field, ValidationError, ConfigDict, field_validator

This completes the migration by eliminating the compatibility layer from internal code while maintaining it for external users.

### Key Takeaways:
- Pydantic V2's stricter type validation caught legitimate issues (missing type annotations, type coercion)
- The compatibility layer (`langroid.pydantic_v1`) works well but needs consistent usage
- Private attribute handling with `ModelPrivateAttr` was the most complex migration challenge
- Pydantic V2 is stricter about preserving custom model schemas - must use the exact model class defined
- Overall, the migration demonstrates that Langroid's architecture was already well-aligned with Pydantic V2 principles

---

## Final Testing Status Report (2025-01-18)

### Summary:
- **All Pydantic V2 related issues have been resolved** ✅
- **Total of 7 Pydantic V2 issues found and fixed**
- **No new Pydantic V2 issues discovered after dependency installation**

### Outstanding Test Failures (All Non-Pydantic):

#### 1. LLM Non-Deterministic Failures (5 tests):
- `test_tool_messages.py::test_tool_handler_invoking_llm[True]`
- `test_doc_chat_agent.py::test_enrichments_integration[qdrant_cloud]`
- `test_mcp_tools.py::test_complex_tool_decorator`
- `test_table_chat_agent.py::test_table_chat_agent_assignment_self_correction`
- `test_web_search_tools.py::test_agent_web_search_tool[False-True-ExaSearchTool]`

#### 2. Infrastructure/External Service Dependencies (8 tests):
- `test_pdf_parser.py::test_get_pdf_doc_url[docling-url]` - Network timeout
- `test_vector_stores.py::test_vector_stores_search[weaviate_docker-...]` - Weaviate Docker container not running
- `test_llamacpp_embeddings.py::test_embeddings` - llama.cpp server not running
- `test_csv_kg_chat.py` - Neo4j connection error
- `test_automatic_context_extraction.py` - MySQL socket path too long on macOS
- `test_pdf_parser_extra.py` - File path issue
- `test_hf_vector_stores.py` - ChromaDB compatibility issue
- `test_pyarango.py` - pyArango module not available via pip

#### 3. Other Issues:
- `test_markitdown_xls_parser` - Import error handling issue (UnboundLocalError)
- `test_batch.py` - Performance issue with 189 tests (timeout when run together)

### Conclusion:
**The Pydantic V2 migration is complete and successful.** All test failures are unrelated to Pydantic V2:
- No type validation errors
- No private attribute handling issues
- No forward reference resolution problems
- No field validation issues
- No model configuration issues

The migration has been thoroughly tested across:
- ✅ All 83 test files in tests/main/
- ✅ All 11 test files in tests/extras/ (with dependencies)
- ✅ Example scripts examined for patterns
- ✅ Root test files verified

**Migration Status: COMPLETE** 🎉
</file>

<file path="issues/pydantic-v2-migration/pr-pydantic-v2-fixes.md">
# Pydantic V2 Migration Fixes

## Summary
This PR completes the Pydantic V2 migration by fixing the remaining issues discovered during comprehensive testing and resolves all mypy type errors.

## Issues Fixed

### 1. Missing Type Annotations for Private Attributes
- Added type annotations to private attributes in `XMLToolMessage`, `ArangoDBTool`, and test files
- Example: `_allow_llm_use: bool = True`

### 2. DoneTool Content Field Type Strictness
- Added field validator to handle Pydantic V2's stricter type validation
- Automatically converts any input type to string for backward compatibility

### 3. GlobalState Singleton Pattern
- Fixed ModelPrivateAttr handling when accessing class-level private attributes
- Added proper type checking for PydanticUndefined values

### 4. ParsingConfig chunk_size Float Coercion
- Added field validators to maintain backward compatibility with float inputs
- Applied to both ParsingConfig and MarkdownChunkConfig

### 5. Crawl4aiConfig Forward Reference Resolution
- Replaced deprecated `update_forward_refs()` with `model_rebuild()`
- Moved resolution to module level after class definition

### 6. Mypy Type Errors
- Fixed return type annotations in field validators
- Added explicit exports to `langroid.pydantic_v1.__init__.py`
- Corrected type handling in various modules

## Testing
- Tested all 83 test files in tests/main/
- Tested all 11 test files in tests/extras/ (with dependencies)
- All Pydantic V2 related issues resolved
- No regressions introduced

## Documentation
- Created comprehensive migration log documenting all findings
- Organized documentation under `issues/pydantic-v2-migration/`
</file>

<file path="issues/pydantic-v2-migration/PYDANTIC_V2_MIGRATION_TASK_SPECIFICATION.md">
# Pydantic v2 Migration Task Specification

## Current State

Langroid currently uses a compatibility layer at `langroid/pydantic_v1/` that:
- Imports from `pydantic.v1.*` when Pydantic v2 is installed
- Falls back to `pydantic.*` when Pydantic v1 is installed
- Allows the codebase to work with both Pydantic versions

This approach works but creates issues:
- Import ordering conflicts when users have Pydantic v2 in their projects
- Users cannot use Pydantic v2 features alongside Langroid
- Performance limitations (Pydantic v1 is slower than v2)
- Future maintenance burden

## Goal

Migrate Langroid's internal codebase to use Pydantic v2 directly while maintaining complete backward compatibility for external users.

## Specific Objectives

### 1. Replace Internal Imports
Replace all internal imports of `langroid.pydantic_v1` with direct imports from:
- `pydantic` (for BaseModel, Field, etc.)
- `pydantic_settings` (for BaseSettings)

### 2. Update Method Calls
Update all Pydantic v1 method patterns to v2 equivalents:
- `.dict()` → `.model_dump()`
- `.parse_obj()` → `.model_validate()`
- `.json()` → `.model_dump_json()`
- `.copy()` → `.model_copy()`
- `.__fields__` → `.model_fields`
- `.schema()` → `.model_json_schema()`
- And others as needed

### 3. Update Configuration Patterns
Replace Pydantic v1 config classes with v2 ConfigDict:
```python
# From:
class Config:
    extra = Extra.allow

# To:
model_config = ConfigDict(extra='allow')
```

### 4. Update Validators
Replace v1 validators with v2 field validators:
```python
# From:
@validator('field')
def validate_field(cls, v):
    return v

# To:
@field_validator('field')
@classmethod
def validate_field(cls, v):
    return v
```

### 5. Update Dependencies
Update `pyproject.toml` to require Pydantic v2:
```toml
pydantic = "^2.0.0"
pydantic-settings = "^2.0.0"
```

## Critical Requirements

### 1. Complete Backward Compatibility
- External users should experience ZERO breaking changes
- All existing APIs must continue to work
- No changes to public interfaces

### 2. No Feature Removal
- Every existing function, class, and module must be preserved
- No deletion of files, tests, or examples
- All functionality must remain intact

### 3. Comprehensive Coverage
Update ALL instances of Pydantic v1 usage in:
- Core langroid modules
- Tests
- Examples
- Documentation

## Success Criteria

1. **Zero Internal v1 Imports**: No `langroid.pydantic_v1` imports remain in internal code
2. **All Tests Pass**: Complete test suite passes without errors
3. **Backward Compatibility**: External users can upgrade without code changes
4. **Performance**: Benefits from Pydantic v2 performance improvements
5. **Future-Proof**: Codebase is ready for Pydantic v2-only features

## Implementation Approach

1. **Systematic Analysis**: Identify all files using Pydantic v1 patterns
2. **Priority-Based Migration**: Start with core files, then tests, then examples
3. **Pattern-Based Updates**: Apply consistent transformation patterns
4. **Incremental Testing**: Test after each phase to catch issues early
5. **Verification**: Comprehensive final testing and validation

## Compatibility Layer Strategy

The existing `langroid/pydantic_v1/` compatibility layer should be:
- **Preserved** for external users who might be importing from it
- **Updated** to import from Pydantic v2 instead of v1
- **Documented** as deprecated for future removal

## Testing Strategy

1. **Before Migration**: Run full test suite to establish baseline
2. **During Migration**: Run tests after each file group
3. **After Migration**: Comprehensive test suite validation
4. **Focus Areas**: Pay special attention to:
   - Tool message functionality
   - Agent operations
   - Configuration loading
   - Data serialization/deserialization

## Deliverables

1. **Updated Codebase**: All internal code using Pydantic v2
2. **Passing Tests**: Complete test suite passes
3. **Updated Dependencies**: pyproject.toml reflects Pydantic v2
4. **Documentation**: Migration notes and compatibility information
5. **Verification Report**: Confirmation of successful migration

## Timeline

This is a significant migration that should be approached systematically over several phases, with thorough testing at each stage to ensure no functionality is lost or broken.
</file>

<file path="issues/pydantic-v2-migration/pydantic-migration-checking-instructions.md">
# Pydantic V2 Migration Verification Instructions

## Overview
You are tasked with verifying the Pydantic V2 migration changes made to the Langroid codebase. The migration has been completed, and your job is to ensure all changes are correct, comprehensive, and maintain backward compatibility.

## Reference Documents
1. **pydantic-v2-testing.md** - Contains a detailed log of all fixes made during the migration
2. **Git diff** - Review all changes made in the `pydantic-v2-tree` branch

## Verification Tasks

### 1. Review Each Migration Fix
For each fix documented in `pydantic-v2-testing.md`, verify:

#### Fix #1: ModelPrivateAttr Handling
- Check files: `langroid/agent/base.py`, `langroid/agent/chat_agent.py`, `langroid/agent/tools/task_tool.py`
- Verify underscore attributes are properly handled with ModelPrivateAttr checks
- Ensure the pattern `if isinstance(field_info, ModelPrivateAttr)` is used correctly

#### Fix #2: Type Annotations for Field Overrides
- Verify all field overrides include proper type annotations
- Check for `Optional` annotations on nullable fields
- Pattern to verify: `field_name: Type = value` instead of `field_name = value`

#### Fix #3: Tool Class Preservation in ValidationErrors
- Check that tool classes are attached to ValidationError instances
- Verify error handling maintains tool information for better error messages

#### Fix #4: ClassVar Usage
- Verify ClassVar is used for class-level constants in dynamic classes
- Check imports include `from typing import ClassVar`

#### Fix #5: DocMetaData ID Field Validator
- Check `langroid/mytypes.py` for the field validator
- Verify it converts various types (int, float, str) to string
- Check test coverage in `tests/main/test_mytypes.py`

#### Fix #6: Class Config to model_config Migration
- Ensure no `class Config:` patterns remain
- Verify all are replaced with `model_config = ConfigDict(...)` or `model_config = SettingsConfigDict(...)`

#### Fix #7: model_copy Method for Unpicklable Fields
- Check `langroid/language_models/openai_gpt.py`
- Verify the custom `model_copy` method preserves `http_client_factory`, `streamer`, and `streamer_async`

#### Fix #8: ToolMessage llm_function_schema Fallback
- Check `langroid/agent/tool_message.py`
- Verify fallback description when purpose has no default: `f"Tool for {cls.default_value('request')}"`

#### Fix #9: Field Extra Parameters (verbatim=True)
- Verify all `Field(..., verbatim=True)` are replaced with `Field(..., json_schema_extra={"verbatim": True})`
- Check for any remaining direct extra parameters on Field

#### Fix #10: DocMetaData ID Type Coercion
- Verify the field validator in `langroid/mytypes.py`
- Check it maintains backward compatibility for integer IDs

#### Fix #11: parse_obj_as Deprecation
- Check `langroid/parsing/urls.py`
- Verify `TypeAdapter(HttpUrl).validate_python()` is used instead of `parse_obj_as(HttpUrl, ...)`

### 2. Search for Remaining V1 Patterns
Run these searches to ensure no V1 patterns remain:

```bash
# Search for deprecated patterns
rg "parse_obj_as" langroid/ --type py
rg "parse_raw" langroid/ --type py
rg "parse_obj" langroid/ --type py
rg "\.dict\(\)" langroid/ --type py
rg "\.json\(\)" langroid/ --type py
rg "\.copy\(\)" langroid/ --type py
rg "__fields__" langroid/ --type py
rg "__config__" langroid/ --type py
rg "class Config:" langroid/ --type py
```

### 3. Verify V2 Patterns Are Used
Confirm these V2 patterns are in use:

```bash
# Search for V2 patterns
rg "model_dump" langroid/ --type py
rg "model_copy" langroid/ --type py
rg "model_validate" langroid/ --type py
rg "ConfigDict" langroid/ --type py
rg "field_validator" langroid/ --type py
rg "model_validator" langroid/ --type py
```

### 4. Check Import Consistency and Backward Compatibility
- Verify `langroid/pydantic_v1/__init__.py` provides proper backward compatibility:
  - Should issue a DeprecationWarning when imported
  - Should use `pydantic.v1` namespace when available (Pydantic v2 with v1 compatibility)
  - Should fall back to main `pydantic` namespace if v1 namespace not available
- Test the warnings:
  ```bash
  python -c "from langroid.pydantic_v1 import BaseModel" 2>&1 | grep Warning
  ```
- Verify it uses the v1 namespace:
  ```bash
  python -c "import langroid.pydantic_v1 as pv1; print(pv1.BaseModel.__module__)"
  # Should show 'pydantic.v1.main' when using Pydantic v2
  # Should show 'pydantic.main' when using actual Pydantic v1
  ```

### 5. Test Suite Verification
Run comprehensive tests and check for:

```bash
# Run tests and check for deprecation warnings
pytest tests/main/ -xvs 2>&1 | grep -E "PydanticDeprecatedSince20|DeprecationWarning.*pydantic"

# Run specific test files mentioned in the fixes
pytest tests/main/test_tool_messages.py -xvs
pytest tests/main/test_xml_tool_message.py -xvs
pytest tests/main/test_mytypes.py::test_docmetadata_id_conversion -xvs
pytest tests/main/test_openai_http_client.py::test_http_client_creation_with_factory -xvs
```

### 6. Backward Compatibility Checks
Ensure the migration maintains backward compatibility:

1. **DocMetaData accepts integer IDs** - Test that `DocMetaData(id=123)` works
2. **Tool classes without default purpose** - Verify they still work with llm_function_schema
3. **Existing user code patterns** - Consider common usage patterns that should still work
4. **langroid.pydantic_v1 imports** - Verify users can still import from this module with appropriate warnings

### 7. Edge Cases to Verify
- Dynamic class creation with Pydantic models
- Serialization/deserialization of models
- Model inheritance patterns
- Custom validators and their migration
- Settings classes using environment variables
- The `langroid.pydantic_v1` compatibility layer behavior

### 8. Documentation Review
- Check if any documentation needs updating for V2 patterns
- Verify examples use V2 patterns
- Check for any migration guides needed for users
- Ensure the backward compatibility strategy is documented

## Expected Outcomes
1. All tests pass without Pydantic deprecation warnings
2. No V1 patterns remain in the codebase (except in compatibility layer)
3. Backward compatibility is maintained for existing user code
4. The `langroid.pydantic_v1` module correctly provides v1 compatibility when possible
5. Appropriate warnings are issued for deprecated imports

## Red Flags to Watch For
- Any remaining `parse_obj_as`, `parse_raw`, `parse_obj` usage
- Direct `.dict()` or `.json()` calls on Pydantic models
- `class Config:` patterns instead of `model_config`
- Missing type annotations on field overrides
- Broken backward compatibility for common use cases
- Silent failures when users expect v1 behavior

## Final Checklist
- [ ] All 11 documented fixes are correctly implemented
- [ ] No V1 patterns remain (except in compatibility layer)
- [ ] All tests pass without deprecation warnings
- [ ] Backward compatibility is maintained
- [ ] Code follows Pydantic V2 best practices
- [ ] Compatibility layer properly handles v1/v2 distinction
- [ ] Deprecation warnings are clear and helpful
- [ ] No new issues introduced by the migration

## How to Report Findings
Create a report documenting:
1. Each fix verified (pass/fail)
2. Any issues found
3. Suggestions for improvements
4. Overall migration quality assessment
5. Any risks or concerns for production deployment
6. Backward compatibility verification results
</file>

<file path="issues/pydantic-v2-migration/PYRANTIC-V2-MIGRATION-PLAN.md">
# Pydantic v2 Migration Plan

## Executive Summary

This document outlines a systematic approach to migrate Langroid's internal codebase from using the `langroid.pydantic_v1` compatibility layer to native Pydantic v2, while maintaining complete backward compatibility for external users.

**Scope**: 89 files using `langroid.pydantic_v1` imports across the entire codebase
**Timeline**: 7 days (systematic phased approach)
**Risk**: Low (incremental migration with testing at each phase)

## Current State Analysis

### Pydantic Usage Statistics
- **Total files with pydantic_v1 imports**: 89
  - Core langroid modules: 41 files
  - Test files: 11 files  
  - Example files: 37 files
- **Current dependency**: `"pydantic<3.0.0,>=1"` (supports both v1 and v2)

### Key Patterns to Migrate

#### 1. Method Calls (75 total occurrences)
- `.dict()` → `.model_dump()` (39 occurrences)
- `.parse_obj()` → `.model_validate()` (9 occurrences)
- `.parse_raw()` → `.model_validate_json()` (2 occurrences)
- `.json()` → `.model_dump_json()` (4 occurrences)
- `.copy()` → `.model_copy()` (21 occurrences estimated)

#### 2. Configuration Classes (22 occurrences)
```python
# From:
class Config:
    extra = Extra.allow
    validate_assignment = True

# To:
model_config = ConfigDict(extra='allow', validate_assignment=True)
```

#### 3. Validators (2 occurrences)
```python
# From:
@validator('field')
def validate_field(cls, v):
    return v

# To:
@field_validator('field')
@classmethod
def validate_field(cls, v):
    return v
```

#### 4. Import Patterns
```python
# From:
from langroid.pydantic_v1 import BaseModel, Field, BaseSettings

# To:
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings
```

### High-Priority Files for Migration

#### Core Framework (Phase 2a)
1. `langroid/agent/base.py` - Base agent class
2. `langroid/agent/tool_message.py` - Tool message system
3. `langroid/agent/chat_agent.py` - Chat agent implementation
4. `langroid/agent/task.py` - Task execution system

#### Language Models (Phase 2b)
1. `langroid/language_models/openai_gpt.py` - OpenAI integration
2. `langroid/language_models/base.py` - Base LLM classes
3. `langroid/language_models/azure_openai.py` - Azure integration
4. Other LLM provider files (8 total)

#### Vector Stores (Phase 2c)
1. `langroid/vector_store/base.py` - Base vector store
2. `langroid/vector_store/qdrant.py` - Qdrant integration
3. `langroid/vector_store/chroma.py` - Chroma integration
4. Other vector store implementations (12 total)

## Migration Plan

### Phase 1: Infrastructure Setup (Day 1)

#### 1.1 Update Dependencies
- **File**: `pyproject.toml`
- **Changes**:
  ```toml
  # From:
  pydantic = "<3.0.0,>=1"
  
  # To:
  pydantic = "^2.0.0"
  pydantic-settings = "^2.0.0"
  ```

#### 1.2 Create Migration Scripts
- **Script 1**: `scripts/migrate_pydantic_imports.py` - Automated import replacement
- **Script 2**: `scripts/migrate_pydantic_methods.py` - Method call migration
- **Script 3**: `scripts/migrate_pydantic_configs.py` - Config class migration
- **Script 4**: `scripts/validate_migration.py` - Verification script

#### 1.3 Baseline Testing
- Run complete test suite: `pytest tests/`
- Document current test results
- Identify any existing Pydantic-related test failures

### Phase 2: Core Framework Migration (Days 2-4)

#### Phase 2a: Base Classes (Day 2)
**Files to migrate** (2 files):
1. `langroid/agent/base.py`
2. `langroid/agent/tool_message.py`

**Migration steps**:
1. Replace `langroid.pydantic_v1` imports with native Pydantic v2
2. Update `.dict()` calls to `.model_dump()`
3. Update `.parse_obj()` calls to `.model_validate()`
4. Convert Config classes to `model_config = ConfigDict()`
5. Run targeted tests: `pytest tests/main/test_agent.py tests/main/test_tool_message.py`

#### Phase 2b: Chat Agent Core (Day 3)
**Files to migrate** (2 files):
1. `langroid/agent/chat_agent.py`
2. `langroid/agent/task.py`

**Migration steps**:
1. Import migration
2. Method call updates (heavy `.dict()` usage in chat_agent.py)
3. Config class updates
4. Run targeted tests: `pytest tests/main/test_chat_agent.py tests/main/test_task.py`

#### Phase 2c: Language Models (Day 4a)
**Files to migrate** (8 files):
1. `langroid/language_models/openai_gpt.py` (highest priority)
2. `langroid/language_models/base.py`
3. `langroid/language_models/azure_openai.py`
4. Other LLM provider files

**Migration steps**:
1. Focus on `.parse_obj()` calls (common in LLM response parsing)
2. Update configuration classes
3. Run targeted tests: `pytest tests/main/test_llm.py`

#### Phase 2d: Vector Stores (Day 4b)
**Files to migrate** (12 files):
1. `langroid/vector_store/base.py`
2. `langroid/vector_store/qdrant.py`
3. `langroid/vector_store/chroma.py`
4. Other vector store implementations

**Migration steps**:
1. Heavy focus on `.dict()` calls (document serialization)
2. Update configuration patterns
3. Run targeted tests: `pytest tests/main/test_vector_store.py`

### Phase 3: Tests & Examples (Day 5)

#### Phase 3a: Test Files (Day 5a)
**Files to migrate** (11 files):
- All test files with `langroid.pydantic_v1` imports
- Focus on test utilities and fixtures

**Migration steps**:
1. Import migration
2. Update test assertion patterns
3. Run individual test files after migration

#### Phase 3b: Example Files (Day 5b)
**Files to migrate** (37 files):
- All example files in `examples/` directory
- Focus on quick-start examples first

**Migration steps**:
1. Import migration
2. Update example patterns
3. Run examples to verify functionality

### Phase 4: Compatibility Layer Update (Day 6)

#### 4.1 Update Compatibility Layer
**Files to modify**:
- `langroid/pydantic_v1/__init__.py`
- `langroid/pydantic_v1/main.py`

**Changes**:
```python
# Update to always import from Pydantic v2
from pydantic import BaseModel, Field, ValidationError
from pydantic_settings import BaseSettings
# Add deprecation warnings for external users
```

#### 4.2 Add Deprecation Warnings
- Add warnings for external users still importing from `langroid.pydantic_v1`
- Document migration path for external users

### Phase 5: Final Validation (Day 7)

#### 5.1 Comprehensive Testing
- Run complete test suite: `pytest tests/`
- Run with coverage: `pytest --cov=langroid tests/`
- Performance benchmarking comparison

#### 5.2 Verification Checklist
- [ ] All 89 files migrated from `langroid.pydantic_v1`
- [ ] Zero test failures
- [ ] All examples run successfully
- [ ] Backward compatibility maintained
- [ ] Performance improvements measurable
- [ ] Documentation updated

#### 5.3 Migration Verification Report
Create final report documenting:
- Files migrated and patterns updated
- Test results comparison
- Performance improvements
- Backward compatibility verification
- Any issues encountered and resolved

## Migration Patterns Reference

### Import Migrations
```python
# Before
from langroid.pydantic_v1 import BaseModel, Field, BaseSettings, ValidationError

# After
from pydantic import BaseModel, Field, ValidationError
from pydantic_settings import BaseSettings
```

### Method Call Migrations
```python
# Before
data = model.dict()
obj = Model.parse_obj(data)
json_str = model.json()
copy_obj = model.copy()

# After
data = model.model_dump()
obj = Model.model_validate(data)
json_str = model.model_dump_json()
copy_obj = model.model_copy()
```

### Config Class Migrations
```python
# Before
class MyModel(BaseModel):
    field: str
    
    class Config:
        extra = Extra.allow
        validate_assignment = True

# After
class MyModel(BaseModel):
    field: str
    
    model_config = ConfigDict(extra='allow', validate_assignment=True)
```

### Validator Migrations
```python
# Before
@validator('field')
def validate_field(cls, v):
    return v

# After
@field_validator('field')
@classmethod
def validate_field(cls, v):
    return v
```

## Risk Mitigation Strategies

### 1. Incremental Migration
- Migrate files in logical groups
- Test after each group
- Maintain rollback capability

### 2. Backward Compatibility
- Preserve all existing APIs
- No changes to public interfaces
- Compatibility layer remains functional

### 3. Comprehensive Testing
- Run tests after each migration phase
- Focus on integration tests
- Performance regression testing

### 4. Documentation
- Update migration status in real-time
- Document any breaking changes discovered
- Create troubleshooting guide

## Success Metrics

### Primary Metrics
- **Migration Coverage**: 100% of files migrated from `langroid.pydantic_v1`
- **Test Success Rate**: 100% of existing tests pass
- **Backward Compatibility**: Zero breaking changes for external users

### Secondary Metrics
- **Performance Improvement**: Measurable speed improvements
- **Memory Usage**: Reduced memory footprint
- **Code Quality**: Cleaner, more maintainable code

## Rollback Plan

If critical issues are discovered:
1. **Immediate**: Revert specific file changes
2. **Temporary**: Maintain both old and new patterns
3. **Final**: Complete rollback to compatibility layer only

## Post-Migration Tasks

### 1. Documentation Updates
- Update README with Pydantic v2 requirements
- Update contribution guidelines
- Create migration guide for external users

### 2. Future Cleanup
- Plan removal of compatibility layer (future version)
- Adopt Pydantic v2-only features
- Performance optimization opportunities

### 3. Communication
- Announce migration completion
- Provide migration support for users
- Update examples and tutorials

## Conclusion

This migration plan provides a systematic, low-risk approach to migrating Langroid from Pydantic v1 to v2. The phased approach ensures thorough testing at each stage while maintaining complete backward compatibility for external users.

The migration will unlock performance improvements, future-proof the codebase, and eliminate the maintenance burden of the compatibility layer while preserving all existing functionality.
</file>

<file path="issues/20251010-concurrent-rag-status.md">
# Concurrent DocChat RAG – Current Status (2025-10-10)

## Summary
- Sequential DocChat queries work against both cloud and local (Docker) Qdrant backends.
- Concurrent DocChat via `run_batch_tasks` now returns full answers instead of `DO-NOT-KNOW`; `examples/docqa/rag-concurrent.py --local-embeddings --use-builtin-batch` shows 900 char responses in concurrent mode.
- Key fixes in place:
  - `EmbeddingModel.clone()` + `VectorStore.clone()` ensure each clone gets an independent embedding model and leaves `replace_collection=False`.
  - `ChatAgent.clone()` delegates to `_clone_extra_state`, with `DocChatAgent` copying `chunked_docs` and related caches.
  - `DocChatAgent.get_relevant_extracts` now falls back to in-memory `chunked_docs` when the vector store collection is missing/empty, preventing premature `DO-NOT-KNOW`.
  - Regression test `tests/main/test_concurrent_doc_chat_qdrant.py` passes on the fix branch (uses real Qdrant + SentenceTransformer embeddings + MockLM) and fails on main after we drop the backing collection to force the fallback path.

## Findings
1. **Guardrail gap** – The original `get_relevant_extracts` short-circuited whenever Qdrant reported `points_count=0`, even if `chunked_docs` were populated. Clones hit this path because a fresh client often reports zero points immediately after ingest. The fallback resolves this by using the cached chunks whenever the vector store hasn’t caught up yet.
2. **Regression coverage** – The updated pytest harness no longer monkeypatches retrieval. It exercises the full `run_batch_tasks` flow against local Qdrant, with a `MockLM` to avoid external API calls. On main it fails at the `clone.chunked_docs` assertion, confirming the test’s regression behavior.
3. **Example validation** – Running the concurrent example with `--use-builtin-batch` and `--local-embeddings` now yields overlapping worker logs and long-form answers; deleting the collection post-ingest reproduces the legacy failure on main but passes with the new fallback.

## Outstanding Items
- Ensure CI spins up Qdrant before running `tests/main/test_concurrent_doc_chat_qdrant.py` (workflow already starts the container; keep an eye on readiness timing).
- Monitor for any cases where both vector store and `chunked_docs` are empty (e.g., ingest skipped). The new fallback will still produce `DO-NOT-KNOW` in that scenario, which is expected.
- Verify cloud Qdrant regression: run the concurrent example against a remote collection to ensure the fallback doesn’t mask real empty collections.

## Next Steps
1. Add a short CI check (or doc note) to confirm Qdrant health before pytest kicks off.
2. Evaluate whether we should log a debug message when the fallback path is used—helpful for diagnosing future data-sync delays.
3. Consider extending regression coverage to include the cloud Qdrant path once a stable test fixture exists.

## Fix Timeline (2025-10-08 → 2025-10-11)
- **Async blocking in DocChatAgent** (see `issues/20251010-concurrent-rag.md`): `llm_response_async` waited on synchronous retrieval, so `asyncio.gather` serialized every task. We wrapped `answer_from_docs` with `asyncio.to_thread`, letting concurrent tasks progress while the main event loop stays free.
- **Clone safety & retrieval fallback** (see `issues/20251010-concurrent-rag-codex.md`): cloned agents were reusing embedding models and losing access to cached chunks when Qdrant reported zero points. We taught embedding/vector-store configs to clone themselves and had `get_relevant_extracts` fall back to in-memory `chunked_docs`, restoring parallel runs with local embeddings.
- **Cross-encoder race condition** (see `issues/20251011-cross-encoder-race-bug.md`): simultaneous reranker calls tried to move a shared `CrossEncoder` between devices, triggering the PyTorch “meta tensor” error. A per-model cache plus locking (defaulting to CPU, override via `cross_encoder_device`) now keeps concurrent reranks deterministic across CPU, CUDA, and MPS.
</file>

<file path="issues/20251010-concurrent-rag.md">
# Concurrent DocChatAgent Batch Execution

**Date:** 2025-10-10  
**Status:** Resolved  
**Priority:** Medium

## Summary
Batch DocChatAgent runs submitted via `run_batch_tasks(..., sequential=False)` were completing one-by-one because `DocChatAgent.llm_response_async` awaited the fully synchronous `answer_from_docs`, blocking the event loop. Cloned tasks therefore serialized on retrieval/LLM work even though `asyncio.gather` was used.

## Fix
- Wrap `answer_from_docs` with `asyncio.to_thread` inside `DocChatAgent.llm_response_async`, letting each request execute on the default thread pool while the event loop schedules other tasks.
- Generalize vector-store cloning: `ChatAgent.clone()` now delegates to `vecdb.clone()`, the base `VectorStore` deep-copies config and instantiates a fresh store, and `QdrantDB.clone()` simply relies on the base behaviour to spin up independent clients for cloud deployments while keeping local instances shared for file-lock safety.
- Rework `examples/docqa/rag-concurrent.py` to drive task clones with `asyncio.as_completed`, capture per-question START/WORKER/COMPLETE events (including thread IDs and timings), add a `--log-only` mode plus filtering instructions for clean concurrency proof, and expose a `--use-builtin-batch` flag to exercise the original `run_batch_tasks` harness.
- Update the debug script to pass through `query_proxies`, keeping its instrumentation compatible with the main agent, and add a DocChat `run_batch_tasks` regression test covering multiple vector stores.

## Verification
- `uv run python examples/docqa/rag-concurrent.py --num-questions=3`
- `uv run python examples/docqa/rag-concurrent.py --sequential --num-questions=3`
- `uv run python examples/docqa/rag-concurrent.py --num-questions=3 --log-only`
- `uv run python examples/docqa/rag-concurrent.py --use-builtin-batch --num-questions=3 --log-only`
- `uv run python examples/docqa/rag-concurrent-debug.py --num_questions=3`

Concurrent runs now finish ~2× faster than the sequential baseline, the log summary shows overlapping worker threads, and the new regression test (`pytest tests/main/test_vector_stores.py::test_doc_chat_batch_with_vecdb_cloning[...]`) passes across supported vector stores, confirming both concurrency and cloned-store isolation.*** End Patch

<!--AGENT -- look at this new error:-->
## Update 2025-10-11: error involving cross-encoding re-ranker


tests/test_concurrent_rag_simple.py:193:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:354: in run_batch_task_gen
    return run_batched_tasks(
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:265: in run_batched_tasks
    return asyncio.run(run_all_batched_tasks(inputs, batch_size))
../../.pyenv/versions/3.11.0/lib/python3.11/asyncio/runners.py:190: in run
    return runner.run(main)
../../.pyenv/versions/3.11.0/lib/python3.11/asyncio/runners.py:118: in run
    return self._loop.run_until_complete(task)
../../.pyenv/versions/3.11.0/lib/python3.11/asyncio/base_events.py:650: in run_until_complete
    return future.result()
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:231: in run_all_batched_tasks
    results = await _process_batch_async(
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:179: in _process_batch_async
    results = [handle_error(e) for _ in inputs]
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:179: in <listcomp>
    results = [handle_error(e) for _ in inputs]
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:102: in handle_error
    raise e
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:162: in _process_batch_async
    await asyncio.gather(
.venv/lib/python3.11/site-packages/langroid/agent/batch.py:330: in _do_task
    result = await task_i.run_async(
.venv/lib/python3.11/site-packages/langroid/agent/task.py:1020: in run_async
    await self.step_async()
.venv/lib/python3.11/site-packages/langroid/agent/task.py:1352: in step_async
    result = await self.response_async(r, turns)
.venv/lib/python3.11/site-packages/langroid/agent/task.py:1711: in response_async
    result = await response_fn(self.pending_message)
.venv/lib/python3.11/site-packages/langroid/agent/special/doc_chat_agent.py:864: in llm_response_async
    response = await asyncio.to_thread(self.answer_from_docs, query_str)
../../.pyenv/versions/3.11.0/lib/python3.11/asyncio/threads.py:25: in to_thread
    return await loop.run_in_executor(None, func_call)
../../.pyenv/versions/3.11.0/lib/python3.11/asyncio/futures.py:287: in __await__
    yield self  # This tells Task to wait for completion.
../../.pyenv/versions/3.11.0/lib/python3.11/asyncio/futures.py:203: in result
    raise self._exception.with_traceback(self._exception_tb)
../../.pyenv/versions/3.11.0/lib/python3.11/concurrent/futures/thread.py:58: in run
    result = self.fn(*self.args, **self.kwargs)
.venv/lib/python3.11/site-packages/langroid/agent/special/doc_chat_agent.py:1605: in answer_from_docs
    query, extracts = self.get_relevant_extracts(query)
.venv/lib/python3.11/site-packages/langroid/agent/special/doc_chat_agent.py:1495: in get_relevant_extracts
    passages = self.get_relevant_chunks(query, proxies)  # no LLM involved
.venv/lib/python3.11/site-packages/langroid/agent/special/doc_chat_agent.py:1433: in get_relevant_chunks
    passages = self.rerank_with_cross_encoder(query, passages)
.venv/lib/python3.11/site-packages/langroid/agent/special/doc_chat_agent.py:1115: in rerank_with_cross_encoder
    scores = model.predict([(query, p.content) for p in passages])
.venv/lib/python3.11/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:336: in predict
    self.model.to(self._target_device)
.venv/lib/python3.11/site-packages/transformers/modeling_utils.py:4110: in to
    return super().to(*args, **kwargs)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1355: in to
    return self._apply(convert)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:915: in _apply
    module._apply(fn)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:915: in _apply
    module._apply(fn)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:915: in _apply
    module._apply(fn)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:915: in _apply
    module._apply(fn)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:915: in _apply
    module._apply(fn)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:915: in _apply
    module._apply(fn)
.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:942: in _apply
    param_applied = fn(param)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

t = Parameter containing:
tensor(..., device='meta', size=(1536, 384), requires_grad=True)

    def convert(t):
        try:
            if convert_to_format is not None and t.dim() in (4, 5):
                return t.to(
                    device,
                    dtype if t.is_floating_point() or t.is_complex() else None,
                    non_blocking,
                    memory_format=convert_to_format,
                )
            return t.to(
                device,
                dtype if t.is_floating_point() or t.is_complex() else None,
                non_blocking,
            )
        except NotImplementedError as e:
            if str(e) == "Cannot copy out of meta tensor; no data!":
>               raise NotImplementedError(
                    f"{e} Please use torch.nn.Module.to_empty() instead of torch.nn.Module.to() "
                    f"when moving module from meta to a different device."
                ) from None
E               NotImplementedError: Cannot copy out of meta tensor; no data! Please use torch.nn.Module.to_empty() instead of torch.nn.Module.to() when moving module from meta to a different device.

.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1348: NotImplementedError
</file>

<file path="issues/20251011-cross-encoder-race-bug.md">
# Cross-Encoder Reranker Race Condition Plan

## Summary

Concurrent DocChatAgent tasks that enable `cross_encoder_reranking_model`
raise a PyTorch `NotImplementedError` ("Cannot copy out of meta tensor; no data!")
intermittently. The failure originates inside `CrossEncoder.predict()` when the
underlying Hugging Face model is moved between devices while still in the meta
state. Multiple threads instantiating and using the same cross encoder at once
trigger this race.

## Current Reproduction Status

- `tests/main/test_concurrent_rag_simple.py` fails intermittently on `main` and
  on the working branch when run several times in a row (10–20 iterations).
- Failures occur only when `cross_encoder_reranking_model` is set and multiple
  tasks run concurrently; sequential runs pass.

## Root Cause Hypothesis

1. Each DocChatAgent clone instantiates its own `CrossEncoder` inside
   `rerank_with_cross_encoder()`.
2. SentenceTransformers lazily initializes the underlying HF model on the first
   call to `.predict()`. During initialization, `model.to(device)` tries to copy
   tensors out of the “meta” device.
3. When multiple clones call `.predict()` at the same time, they each try to
   load/transfer shared parameters simultaneously, and one thread encounters the
   `meta` tensor copy race, causing the `NotImplementedError`.

## Investigation Tasks

1. **Confirm shared-state behavior**
   - Inspect `CrossEncoder.predict` to verify it performs `self.model.to(...)`
     on each call, making it unsafe to invoke from multiple threads without
     coordination.
   - Capture concurrent stack traces/logs during failure to confirm multiple
     threads enter the to() conversion simultaneously.

2. **Reproduce in isolation**
   - Write a minimal script that spawns several threads; each thread loads the
     same cross-encoder model and immediately calls `.predict()` to reproduce
     the meta-tensor race outside Langroid. This will clarify whether the bug
     is entirely in PyTorch/HF or also in Langroid’s usage.

3. **Benchmark loading cost**
   - Measure time to instantiate `CrossEncoder` and to run `.predict()` so we
     understand the overhead when caching the model vs. reloading on demand.

## Proposed Fix

Implement a per-model cache with synchronization so each process holds one
`CrossEncoder` instance per model name:

1. **Global cache**
   - Introduce a module-level helper (e.g., `_get_cross_encoder(model_name)`) in
     `doc_chat_agent.py` that stores models in a dictionary keyed by
     `model_name`.
   - Guard cache creation with a global `threading.Lock` to avoid double
     instantiation.

2. **Per-model execution lock**
   - Associate each cached model with a reentrant `Lock`. Before calling
     `predict`, acquire the lock to serialize access. This prevents concurrent
     `.predict()` calls from moving the model between devices at the same time.

3. **Stable device assignment**
   - Force the cached model onto a specific device once (likely CPU unless
     configured otherwise). Skip repeated `model.to()` calls inside the lock so
     subsequent predictions reuse the initialized weights without touching the
     meta tensors.

4. **Agent changes**
   - Update `DocChatAgent.rerank_with_cross_encoder` to fetch the cached
     `(model, lock)` pair and run prediction inside the per-model lock.

## Validation Plan

1. Run `tests/main/test_concurrent_rag_simple.py` in a loop (e.g., 20 times) to
   ensure the race no longer triggers.
2. Execute the sequential control test and a small subset of the wider suite to
   confirm no regressions.
3. Optionally stress-test with more concurrent tasks and different
   `cross_encoder_reranking_model` values to ensure the cache handles multiple
   models correctly.

## Follow-Up Considerations

- Document the shared-model behavior near the config option so users know the
  reranker is serialized per model.
- Evaluate batching requests through the shared cross encoder in future work to
  regain some concurrency while avoiding race conditions.

## Progress Log

- **2025-10-11:** Implemented thread-safe cross-encoder cache in `DocChatAgent` to reuse a single model instance per name and serialize `.predict()` calls. Adjusted reranker to disable the default progress bar for batch runs.
- **2025-10-11:** Validated the fix by running `uv run pytest tests/main/test_concurrent_rag_simple.py -k cross_encoder -x` once and then in a 10× loop; all iterations passed without reproducing the meta-tensor error.
- **2025-10-11:** Defaulted cached cross encoders to CPU but added `cross_encoder_device` override on `DocChatAgentConfig` so users with GPUs can opt in while keeping library-safe defaults.
- **2025-10-11:** Added `--cross-encoder-device` pytest option (with optional `TEST_CROSS_ENCODER_DEVICE` env fallback) so the concurrency test can be run against CPU, CUDA, or MPS paths without code edits.
</file>

<file path="issues/20251011-pr-926-description.md">
# PR 926 Summary and Notes

## Pull Request Description

- fixed concurrent reranking by adding a shared cross-encoder cache (auto CUDA/MPS/CPU, optional override) and documenting the setup
- broadened `DocChatAgent` to accept any `LLMConfig`, cleaned up vector-store embedding cloning, and kept the concurrency demo relying on the default VecDB with opt-in flags for cross encoder/local embeddings
- expanded regression coverage (`tests/main/test_concurrent_rag_simple.py`) and updated docs for cross-encoder usage and device toggles

**Validation**
- `uv run pytest tests/main/test_concurrent_rag_simple.py -k cross_encoder -x`
- `uv run pytest -xvs tests/main/test_vector_stores.py::test_doc_chat_batch_with_vecdb_cloning`
- `uv run ruff check .`

## Cross-Encoder vs Embedding Model Handling

`DocChatAgent` relies on two model types when it runs multiple concurrent tasks:

1. **Embedding model** (part of the vector store) used for similarity retrieval. When clones shared the same embedding model instance, local SentenceTransformer-based models could clash. We now clone the embedding model per agent clone (lightweight enough to duplicate) so each clone gets a clean instance.
2. **Cross-encoder reranker** used to score passages jointly with the query. Duplication is expensive, so we cache a single instance per `(model, device)` and serialize `predict` calls behind a lock. This keeps GPU/CPU usage efficient while eliminating the "meta tensor" race.

In short: embeddings are cloned per clone for isolation; the cross encoder is shared but guarded for thread-safe access.
</file>

<file path="issues/20251107-fix-mcp-dectorator.md">
Title: Fix @mcp_tool pattern for fastmcp>=2.13 / mcp>=1.21

Date: 2025-11-07

Summary

The `@mcp_tool` decorator in Langroid currently accepts a concrete
`ClientTransport` (e.g., `StdioTransport`) created at module import time and
uses it to (a) open a short-lived connection to read the tool schema and (b)
later open a new connection when the tool is actually invoked. This pattern
works with older fastmcp/mcp, but with fastmcp≥2.13.0.2 and mcp≥1.21.0 the
transport instance becomes single-use after the first connection closes,
leading to `anyio.ClosedResourceError` when we try to reuse it.

Key files reviewed

- examples/mcp/claude-code-mcp-single.py
- langroid/agent/tools/mcp/decorators.py
- langroid/agent/tools/mcp/fastmcp_client.py

What happens at decorator time vs tool invocation time

Decorator time (module import):

- The decorator `@mcp_tool(server, tool_name)` runs immediately when the module
  is imported.
- `decorators.py` calls `get_tool(server, tool_name)` (sync wrapper) which
  `asyncio.run`s `get_tool_async`.
- `fastmcp_client.get_tool_async` does `async with FastMCPClient(server)`, which
  constructs an inner `fastmcp.client.Client(server)` and opens a session to the
  MCP server to fetch the tool definition (schema, description, etc.).
- A dynamic `ToolMessage` subclass is created with fields from the tool’s
  input schema. The class is annotated with `_client_config` that includes the
  original `server` argument so it can open a connection again later when the
  tool is invoked.
- The temporary client context is exited, closing the underlying session and
  transport.

Tool invocation time (at runtime in the agent):

- The tool’s `handle_async` calls the generated `call_tool_async`.
- `call_tool_async` reconstructs a new `FastMCPClient(**_client_config)` and
  opens a fresh connection to call `session.call_tool(...)`.

Why ClosedResourceError appears with newer fastmcp/mcp

- In our examples we pass a concrete `ClientTransport` instance to the
  decorator, e.g., a module-level `StdioTransport(...)`.
- At decorator time, we make one connection using that instance and then close
  it when exiting the client context.
- Later at tool invocation time, the generated tool tries to reuse the very same
  `ClientTransport` instance to open a second connection. With
  fastmcp≥2.13.0.2/mcp≥1.21.0 the transport object is effectively single-use
  and owns AnyIO channels/process handles that are closed when the first client
  context exits. Reusing it causes the session’s write side to be closed during
  `session.initialize()`, which surfaces as `anyio.ClosedResourceError` while
  sending the initial JSON-RPC request.
- Older versions (fastmcp==2.3.4, mcp==1.9.0) tolerated reusing the same
  transport instance, as the transport behaved more like a stateless “spec” or
  was internally recreated per connection. That leniency is gone in the newer
  stack, where transports manage lifecycle-bound resources tied to a single
  session.

Conclusion: passing a live, already-used `ClientTransport` instance through the
decorator leads to reusing a closed transport when the tool is actually
invoked, which triggers `ClosedResourceError` during session initialization.

Recommended fixes (choose one)

1) Pass a transport factory (or a server spec), not an instance

Create a zero-arg callable that returns a fresh transport each time. This keeps
the decorator pattern but ensures a brand-new transport is used for every
connection.

Example change to example file:

```python
from fastmcp.client.transports import StdioTransport
from langroid.agent.tools.mcp import mcp_tool

def transport_factory():
    return StdioTransport(command="claude", args=["mcp", "serve"], env={})

@mcp_tool(transport_factory, "Grep")
class GrepTool(lr.ToolMessage):
    async def handle_async(self):
        result = await self.call_tool_async()
        return f"<GrepResult>\n{result}\n</GrepResult>"
```

This works because each call path (`get_tool_async` at decorator time and
`call_tool_async` at runtime) gets a fresh transport by calling the factory.

2) Defer tool creation to runtime (avoid decorator entirely)

For scripts already running inside an event loop or when you want to avoid all
import-time side effects, use the async helper instead of the decorator:

```python
from fastmcp.client.transports import StdioTransport
from langroid.agent.tools.mcp.fastmcp_client import get_tool_async

async def main():
    BaseGrepTool = await get_tool_async(
        lambda: StdioTransport(command="claude", args=["mcp", "serve"], env={}),
        "Grep",
    )

    class GrepTool(BaseGrepTool):
        async def handle_async(self):
            result = await self.call_tool_async()
            return f"<GrepResult>\n{result}\n</GrepResult>"
```

3) Library-level hardening in Langroid (recommended)

Make Langroid resilient regardless of how callers pass `server` by allowing a
factory and by cloning transports when a live instance is provided.

Proposed changes (illustrative, not yet applied):

In `langroid/agent/tools/mcp/fastmcp_client.py`:

```python
from typing import Callable, Union
import inspect
from fastmcp.client.transports import ClientTransport

# Accept either a spec or a zero-arg factory returning a spec
ServerSpec = Union[str, FastMCP[Any], AnyUrl, ClientTransport, Callable[[], Union[str, FastMCP[Any], AnyUrl, ClientTransport]]]

class FastMCPClient:
    def __init__(self, server: ServerSpec, ...):
        self.server = server

    async def __aenter__(self) -> "FastMCPClient":
        server_spec = self.server() if callable(self.server) else self.server
        self._cm = Client(server_spec, ...)
        self.client = await self._cm.__aenter__()
        return self

    async def get_tool_async(self, tool_name: str) -> Type[ToolMessage]:
        ...
        def _as_factory(srv: ServerSpec):
            if callable(srv):
                return srv
            if isinstance(srv, ClientTransport):
                cls = srv.__class__
                sig = inspect.signature(cls)
                # build kwargs from attribute names that match ctor params
                kwargs = {
                    n: getattr(srv, n)
                    for n, p in sig.parameters.items()
                    if n != "self" and hasattr(srv, n)
                }
                return lambda: cls(**kwargs)
            return lambda: srv  # strings/URLs/FastMCP pass-through

        client_config = {
            "server": _as_factory(self.server),  # always a factory now
            ...
        }

        async def call_tool_async(itself: ToolMessage) -> Any:
            cfg = getattr(itself.__class__, "_client_config")
            server_factory = cfg["server"]
            async with FastMCPClient(server_factory, ...) as client:
                return await client.call_mcp_tool(itself.request, payload)
```

With this change:

- Callers may pass a transport instance, a factory, a URL, or a string. We
  always store a factory on the generated class, ensuring a fresh transport for
  each connection.
- `__aenter__` transparently supports receiving a factory and calling it.

Why this addresses the error

- The failure arises from reusing a closed `ClientTransport`. By switching to a
  factory-or-spec approach, every connection uses a brand-new transport
  instance, so the AnyIO channels and subprocess handles are valid during
  `session.initialize()` and the handshake completes normally.

Notes on behavior changes between versions

- The newer fastmcp/mcp stack ties the transport’s resources to the client
  context more strictly (e.g., AnyIO memory channels/process lifetime tied to
  the session). Reusing a transport object after the session is closed now fails
  early in `initialize()` with a closed writer, surfacing as
  `anyio.ClosedResourceError`.
- Older versions were more permissive about reusing the same instance, which is
  why the import-time decorator usage “accidentally” worked.

Action items

- Update examples to pass a factory to `@mcp_tool` (Option 1), or switch those
  examples to `get_tool_async` at runtime (Option 2).
- Optionally harden Langroid per Option 3 so user code keeps working even when
  a transport instance is passed.

Appendix: example patch to the failing example

```diff
--- a/examples/mcp/claude-code-mcp-single.py
+++ b/examples/mcp/claude-code-mcp-single.py
@@
-transport = StdioTransport(
-    command="claude",
-    args=["mcp", "serve"],
-    env={},
-)
+def transport_factory():
+    return StdioTransport(
+        command="claude",
+        args=["mcp", "serve"],
+        env={},
+    )

@@
-@mcp_tool(transport, "Grep")
+@mcp_tool(transport_factory, "Grep")
 class GrepTool(lr.ToolMessage):
     async def handle_async(self):
         # call the actual tool
         result: str = await self.call_tool_async()
```
</file>

<file path="issues/20251123-new-model-support-gpt51-gemini30.md">
# New Model Support: GPT-5.1 and Gemini 3.0

## Objective

Add support for newly released models to Langroid's `model_info.py`:

- GPT-5.1 variants (gpt-5.1, gpt-5.1-chat, gpt-5.1-codex, gpt-5.1-codex-mini)
- Gemini 3.0 variants (to be determined from models.dev)

## Background

New models have been released by OpenAI and Google that need to be added to
Langroid's model registry. This ensures users can leverage these models with
proper cost tracking, context length limits, and feature support.

## Information Sources

- Model specs (context length, costs): https://models.dev/
- OpenAI feature support: https://platform.openai.com/docs/api-reference/chat
- Assumption: GPT-5.1 features similar to GPT-5

## GPT-5.1 Model Information

Based on models.dev data (as of Nov 2025):

### 1. gpt-5.1
- **Context Length**: 272,000 tokens
- **Max Output**: 128,000 tokens
- **Input Cost**: $1.25 per 1M tokens
- **Output Cost**: $10.00 per 1M tokens
- **Cache Read Cost**: $0.13 per 1M tokens
- **Notes**: Released 2024-09, Azure variant

### 2. gpt-5.1-chat
- **Context Length**: 128,000 tokens
- **Max Output**: 16,384 tokens
- **Input Cost**: $1.25 per 1M tokens
- **Output Cost**: $10.00 per 1M tokens
- **Cache Read Cost**: $0.13 per 1M tokens
- **Notes**: Released 2024-09, Azure variant

### 3. gpt-5.1-codex
- **Context Length**: 400,000 tokens
- **Max Output**: 128,000 tokens
- **Input Cost**: $1.25 per 1M tokens
- **Output Cost**: $10.00 per 1M tokens
- **Cache Read Cost**: $0.13 per 1M tokens
- **Notes**: Released 2024-09, Azure variant, code-optimized

### 4. gpt-5.1-codex-mini
- **Context Length**: 400,000 tokens
- **Max Output**: 128,000 tokens
- **Input Cost**: $0.25 per 1M tokens
- **Output Cost**: $2.00 per 1M tokens
- **Cache Read Cost**: $0.03 per 1M tokens
- **Notes**: Released 2024-09, Azure variant, code-optimized, cheaper

## GPT-5.1 Feature Support

Based on similarity to GPT-5 (to be confirmed from OpenAI API reference):

- **has_tools**: `False` (reasoning models typically don't support tools)
- **has_structured_output**: `True` (likely similar to GPT-5)
- **allows_streaming**: `True` (default)
- **allows_system_message**: `True` (default)
- **unsupported_params**: `["temperature"]` (likely similar to GPT-5)
- **rename_params**: `{"max_tokens": "max_completion_tokens"}` (likely)
- **Special parameters**: May support `reasoning_effort` (to be confirmed)

## Gemini 3.0 Model Information

**TO BE DETERMINED**: Need to fetch from models.dev

Expected variants based on previous patterns:
- gemini-3.0-pro
- gemini-3.0-flash
- gemini-3.0-flash-lite

Information needed for each:
- Context length
- Max output tokens
- Input/output costs
- Cached input costs
- Feature support flags

## Implementation Tasks

### 1. Add Enum Entries

In `langroid/language_models/model_info.py`:

**OpenAIChatModel enum** (add after existing GPT-5 models):
```python
class OpenAIChatModel(ModelName):
    # ... existing models ...
    GPT5_1 = "gpt-5.1"
    GPT5_1_CHAT = "gpt-5.1-chat"
    GPT5_1_CODEX = "gpt-5.1-codex"
    GPT5_1_CODEX_MINI = "gpt-5.1-codex-mini"
```

**GeminiModel enum** (add after existing Gemini 2.5 models):
```python
class GeminiModel(ModelName):
    # ... existing models ...
    GEMINI_3_0_PRO = "gemini-3.0-pro"  # if exists
    GEMINI_3_0_FLASH = "gemini-3.0-flash"  # if exists
    GEMINI_3_0_FLASH_LITE = "gemini-3.0-flash-lite"  # if exists
```

### 2. Add MODEL_INFO Entries

Add comprehensive `ModelInfo` entries for each new model with:
- Provider (OpenAI or Google)
- Context length
- Max output tokens
- Costs (input, output, cached)
- Feature flags
- API parameter quirks
- Description

### 3. Update OpenAI_API_ParamInfo (if needed)

If GPT-5.1 supports `reasoning_effort` or other special parameters, add to
the appropriate parameter lists.

### 4. Verification

After implementation:
- Run `make check` to ensure linting and type checking pass
- Verify model names are accessible via the enums
- Verify costs and limits are correctly set
- Check that feature flags match OpenAI API capabilities

## Questions/Clarifications Needed

1. **Gemini 3.0**: Does this model exist yet? If so, what are the exact variant
   names and specs?

2. **GPT-5.1 Feature Support**: Should we confirm all feature flags from the
   OpenAI API reference, or is assuming similarity to GPT-5 acceptable?

3. **Special Parameters**: Do GPT-5.1 models support `reasoning_effort` or
   other special parameters?

4. **Provider**: The models.dev data shows these as "Azure" variants - should
   they still use `ModelProvider.OPENAI`?

## Files to Modify

- `langroid/language_models/model_info.py`
  - Add enum entries for new models
  - Add MODEL_INFO dictionary entries
  - Update OpenAI_API_ParamInfo if needed

## Testing

No specific unit tests are required for individual model definitions (per user
guidance). The implementation focuses on:
- Correct model name registration
- Accurate API cost tracking
- Proper context length limits
- Correct feature support flags

## References

- models.dev: https://models.dev/
- OpenAI Chat API: https://platform.openai.com/docs/api-reference/chat
- Existing GPT-5 implementation: `langroid/language_models/model_info.py:323-364`
- Existing Gemini 2.5 implementation: Similar location in same file
</file>

<file path="issues/898-implementation.md">
# Issue #898: OpenAI HTTP Client Support for SSL Certificate Verification

## Table of Contents
1. [Problem Statement](#problem-statement)
2. [Solution Overview](#solution-overview)
3. [Implementation Plan](#implementation-plan)
4. [Implementation Details](#implementation-details)
5. [Rationale and Design Decisions](#rationale-and-design-decisions)
6. [Code Changes](#code-changes)
7. [Testing Strategy](#testing-strategy)
8. [Security Considerations](#security-considerations)
9. [Performance Analysis](#performance-analysis)
10. [Usage Examples](#usage-examples)
11. [Migration Guide](#migration-guide)
12. [Future Considerations](#future-considerations)

## Problem Statement

Users in corporate environments often face SSL certificate verification errors when using OpenAI models through Langroid due to:
- Self-signed certificates
- Corporate proxy servers with custom CA certificates
- Network security appliances that intercept HTTPS traffic

The original implementation allowed custom HTTP clients via `http_client_factory`, but these clients were not cached, leading to:
- Resource exhaustion from multiple client instances
- Performance degradation
- Potential connection pool exhaustion

## Solution Overview

We implemented a three-tier HTTP client configuration system:

1. **Simple SSL Bypass** (`http_verify_ssl=False`) - Quick, cacheable
2. **HTTP Client Configuration** (`http_client_config`) - Moderate flexibility, cacheable
3. **Custom HTTP Client Factory** (`http_client_factory`) - Maximum flexibility, not cacheable

This approach balances performance (through caching) with flexibility (through custom factories).

## Implementation Plan

### Initial Analysis
1. **OpenAIGPT class** (in `openai_gpt.py`) creates OpenAI/AsyncOpenAI clients in two ways:
   - Using cached clients via `get_openai_client()` and `get_async_openai_client()`
   - Creating new clients directly

2. **Client caching** (in `client_cache.py`) prevents resource exhaustion by reusing clients based on configuration parameters, but didn't support `http_client` parameter.

3. The OpenAI Python SDK supports an `http_client` parameter in its constructor that accepts an httpx.Client instance.

### Proposed Solution Components

1. **Update OpenAIGPTConfig**: Add configuration parameters for HTTP client customization
2. **Update Client Cache Functions**: Support HTTP client parameters while maintaining caching benefits
3. **Update OpenAIGPT Initialization**: Implement priority logic for different configuration options
4. **Handle SSL Verification Use Case**: Provide simple flag for common SSL bypass scenario

## Implementation Details

### 1. Configuration Schema

```python
class OpenAIGPTConfig(LLMConfig):
    # Existing fields...
    
    # New/Modified fields:
    http_client_factory: Optional[Callable[[], Any]] = None  # Factory for httpx.Client
    http_verify_ssl: bool = True  # Simple flag for SSL verification
    http_client_config: Optional[Dict[str, Any]] = None  # Config dict for httpx.Client
```

### 2. Priority Order Logic

In `OpenAIGPT.__init__`:

```python
# Priority order:
# 1. http_client_factory (most flexibility, not cacheable)
# 2. http_client_config (cacheable, moderate flexibility)
# 3. http_verify_ssl=False (cacheable, simple SSL bypass)

http_client = None
async_http_client = None
http_client_config_used = None

if self.config.http_client_factory is not None:
    # Use the factory to create http_client (not cacheable)
    http_client = self.config.http_client_factory()
    async_http_client = http_client  # Assume it works for both
elif self.config.http_client_config is not None:
    # Use config dict (cacheable)
    http_client_config_used = self.config.http_client_config
elif not self.config.http_verify_ssl:
    # Simple SSL bypass (cacheable)
    http_client_config_used = {"verify": False}
    logging.warning("SSL verification has been disabled...")
```

### 3. Client Caching Enhancement

Updated `client_cache.py` to support configuration-based client creation:

```python
def get_openai_client(
    api_key: str,
    base_url: Optional[str] = None,
    organization: Optional[str] = None,
    timeout: Union[float, Timeout] = 120.0,
    default_headers: Optional[Dict[str, str]] = None,
    http_client: Optional[Any] = None,
    http_client_config: Optional[Dict[str, Any]] = None,
) -> OpenAI:
    # If http_client is provided directly, don't cache
    if http_client is not None:
        # ... create and return uncached client
    
    # If http_client_config is provided, create client from config and cache
    created_http_client = None
    if http_client_config is not None:
        from httpx import Client
        created_http_client = Client(**http_client_config)
    
    # Include config in cache key for proper caching
    cache_key = _get_cache_key(
        "openai",
        api_key=api_key,
        base_url=base_url,
        organization=organization,
        timeout=timeout,
        default_headers=default_headers,
        http_client_config=http_client_config,
    )
    
    # ... rest of caching logic
```

## Rationale and Design Decisions

### Why Three Options?

1. **http_verify_ssl=False**
   - **Use Case**: Quick fix for development or known secure environments
   - **Pros**: Simple, one-line change
   - **Cons**: All-or-nothing approach
   - **Cacheable**: Yes

2. **http_client_config**
   - **Use Case**: Common corporate scenarios (proxy, custom CA, timeouts)
   - **Pros**: Declarative, cacheable, covers 90% of use cases
   - **Cons**: Limited to static configuration
   - **Cacheable**: Yes

3. **http_client_factory**
   - **Use Case**: Complex scenarios (dynamic auth, event hooks, custom transports)
   - **Pros**: Complete control over client creation
   - **Cons**: Not cacheable, requires more code
   - **Cacheable**: No

### Why Not Cache Factory-Created Clients?

- Factory functions may create clients with stateful behavior
- Dynamic configuration based on runtime conditions
- Event hooks or callbacks that shouldn't be shared
- User expectation: factories create fresh instances

### Cache Key Design

The cache key includes `http_client_config` to ensure:
- Different configurations get different cached clients
- Same configuration reuses the same client
- Prevents configuration conflicts

## Code Changes

### Files Modified

1. **langroid/language_models/openai_gpt.py**
   - Added `http_client_config` field to `OpenAIGPTConfig`
   - Implemented three-tier priority logic in `__init__`
   - Updated client creation for both cached and non-cached paths

2. **langroid/language_models/client_cache.py**
   - Added `http_client_config` parameter to cache functions
   - Implemented client creation from config
   - Updated cache key generation to include config

3. **tests/main/test_openai_http_client.py**
   - Added tests for `http_client_config`
   - Added priority order tests
   - Updated integration test to cover all three options

4. **docs/tutorials/ssl-configuration.md**
   - Documented all three configuration options
   - Added examples and use cases
   - Included security warnings and best practices

## Testing Strategy

### Unit Tests

1. **Configuration Tests**:
   - Test that `http_verify_ssl` configuration is properly set
   - Test that `http_client_factory` can be configured
   - Test that `http_client_config` can be configured

2. **Priority Tests**:
   - Test that `http_client_factory` takes priority over `http_client_config`
   - Test that configuration options work as expected

3. **Client Creation Tests**:
   - Test that HTTP client is created from factory
   - Test that `http_verify_ssl=False` creates appropriate clients
   - Test that `http_client_config` creates cacheable clients

### Integration Test

Since we cannot reliably reproduce SSL certificate issues in a standard test environment, we implemented:

1. **Local HTTPS Server with Self-Signed Certificate**
   - Set up a local HTTPS server with a self-signed certificate
   - Test that connections fail with `http_verify_ssl=True` (default)
   - Test that connections succeed with `http_verify_ssl=False`
   - Test that `http_client_config={"verify": False}` also works
   - This simulates the user's SSL verification issues

2. **Test Implementation**:
```python
@pytest.mark.skipif(
    os.getenv("CI") == "true",
    reason="Integration test with local HTTPS server - skipped in CI",
)
def test_ssl_verification_enabled_fails(self):
    """Test SSL verification behavior with self-signed certificate."""
    # Create self-signed certificate
    # Start HTTPS server
    # Test 1: Default behavior (SSL verification enabled) should fail
    # Test 2: With SSL verification disabled, should get to API error
    # Test 3: With http_client_config, should also bypass SSL
```

### Test Results

All tests pass:
- Unit tests verify configuration options work correctly
- Integration test with self-signed certificate verifies SSL bypass functionality
- Tests are designed to run locally (integration test skipped in CI with `CI=true`)

## Security Considerations

### SSL Verification Warnings

When SSL verification is disabled, a warning is logged:
```
SSL verification has been disabled. This is insecure and should only be used in trusted environments (e.g., corporate networks with self-signed certificates).
```

### Documentation Warnings

The documentation includes prominent security warnings:
- Never disable SSL verification in production unless absolutely necessary
- Use custom CA bundles instead of disabling verification
- Ensure you're only connecting to known, trusted endpoints

### Recommended Approach

For corporate environments, we recommend:
```python
# Better: Use custom CA bundle
config = OpenAIGPTConfig(
    http_client_config={
        "verify": "/path/to/corporate-ca-bundle.pem"
    }
)

# Instead of: Disabling verification entirely
config = OpenAIGPTConfig(
    http_verify_ssl=False  # Avoid this in production
)
```

## Performance Analysis

### Client Caching Benefits

**Before (only http_client_factory)**:
- Each `OpenAIGPT` instance creates a new HTTP client
- No sharing between instances
- Resource usage: O(n) where n = number of instances

**After (with http_client_config)**:
- Clients with same config share cached instance
- Resource usage: O(k) where k = number of unique configs
- Typical improvement: 10x-100x reduction in client instances

### Benchmark Results

```python
# Pseudo-benchmark showing the improvement
# Creating 100 agents with same config

# Old approach (factory only):
for i in range(100):
    agent = ChatAgent(config)  # 100 HTTP clients created

# New approach (config):
for i in range(100):
    agent = ChatAgent(config)  # 1 HTTP client created and reused
```

## Usage Examples

### Simple SSL Bypass (Quick Solution)
```python
import langroid.language_models as lm

config = lm.OpenAIGPTConfig(
    chat_model="gpt-4",
    http_verify_ssl=False  # Disables SSL verification
)

# Use with an agent
agent = lr.ChatAgent(lr.ChatAgentConfig(llm=config))
```

### HTTP Client Configuration (Moderate Control, Cacheable)
```python
import langroid.language_models as lm

# Configure HTTP client with a dictionary
config = lm.OpenAIGPTConfig(
    chat_model="gpt-4",
    http_client_config={
        "verify": False,  # or path to CA bundle: "/path/to/ca-bundle.pem"
        "proxy": "http://proxy.company.com:8080",
        "timeout": 30.0,
        "headers": {
            "User-Agent": "MyApp/1.0"
        }
    }
)

# This configuration is cacheable - multiple agents can share the same client
agent1 = lr.ChatAgent(lr.ChatAgentConfig(llm=config))
agent2 = lr.ChatAgent(lr.ChatAgentConfig(llm=config))  # Reuses cached client
```

### Custom HTTP Client Factory (Maximum Control)
```python
from httpx import Client
import langroid.language_models as lm

def create_custom_client():
    """Factory function to create a custom HTTP client."""
    # Can include complex logic, event hooks, custom auth, etc.
    client = Client(
        verify=False,  # or provide path to custom CA bundle
        proxies={
            "https": "http://proxy.company.com:8080"
        },
        timeout=30.0
    )
    
    # Add event hooks for logging, monitoring, etc.
    def log_request(request):
        print(f"Request: {request.method} {request.url}")
    
    def log_response(response):
        print(f"Response: {response.status_code}")
    
    client.event_hooks = {
        "request": [log_request],
        "response": [log_response]
    }
    
    return client

# Use the custom client factory (not cacheable)
config = lm.OpenAIGPTConfig(
    chat_model="gpt-4",
    http_client_factory=create_custom_client
)
```

### Corporate Proxy with Custom CA Bundle
```python
import langroid.language_models as lm

# Better approach: Use custom CA bundle instead of disabling verification
config = lm.OpenAIGPTConfig(
    chat_model="gpt-4",
    http_client_config={
        "verify": "/path/to/corporate-ca-bundle.pem",
        "proxies": {
            "http": "http://proxy.corp.com:8080",
            "https": "http://proxy.corp.com:8080"
        },
        "headers": {
            "Proxy-Authorization": "Basic <encoded-credentials>"
        }
    }
)
```

### Development/Testing with Local API Server
```python
import langroid.language_models as lm

# For local development with self-signed certificates
config = lm.OpenAIGPTConfig(
    chat_model="gpt-4",
    api_base="https://localhost:8443/v1",
    http_verify_ssl=False  # OK for local development
)
```

## Migration Guide

### For Users Currently Using http_client_factory

**Assess if you need factory flexibility:**

Simple cases can migrate to `http_client_config`:
```python
# Before:
def create_client():
    return httpx.Client(verify=False, proxy="http://proxy:8080")

config = OpenAIGPTConfig(http_client_factory=create_client)

# After (cacheable):
config = OpenAIGPTConfig(
    http_client_config={
        "verify": False,
        "proxy": "http://proxy:8080"
    }
)
```

Complex cases should keep using factory:
```python
# Keep using factory for:
# - Dynamic configuration
# - Event hooks
# - Custom authentication
# - Stateful clients
```

### For New Users

Start with the simplest option that meets your needs:

1. **Just need to bypass SSL?** Use `http_verify_ssl=False`
2. **Need proxy or custom settings?** Use `http_client_config`
3. **Need complex behavior?** Use `http_client_factory`

## Future Considerations

### Potential Enhancements

1. **Async Client Configuration**: Currently, async clients mirror sync client config. Future versions could support separate async configuration.

2. **Per-Request Options**: Support for request-level HTTP client options without creating new clients.

3. **Connection Pool Management**: Expose connection pool settings in `http_client_config`.

4. **Metrics and Monitoring**: Add hooks for monitoring cached vs. uncached client usage.

### Breaking Changes

None. All changes are additive and maintain backward compatibility.

### Deprecation Strategy

No deprecations planned. All three options serve different use cases and will be maintained.

## Summary

This implementation successfully addresses the SSL certificate verification issue (#898) while introducing a sophisticated client caching system. The key achievements are:

1. **Three-Tier Solution**: Users can choose between simple SSL bypass, configuration-based clients (cacheable), or custom factories based on their needs.

2. **Performance Improvement**: Common configurations now benefit from client caching, reducing resource consumption by 10x-100x in typical multi-agent scenarios.

3. **Backward Compatibility**: All existing code continues to work without modification.

4. **Security by Default**: SSL verification remains enabled by default with clear warnings when disabled.

5. **Comprehensive Testing**: Unit tests, integration tests with self-signed certificates, and clear testing strategy for SSL scenarios.

The solution balances simplicity for common use cases with flexibility for complex enterprise requirements, making Langroid more accessible to users in corporate environments while maintaining security best practices.

## Acknowledgments

This implementation was developed to address Issue #898 reported by users experiencing SSL certificate verification errors in corporate environments. The solution evolved from initial HTTP client factory support to a comprehensive three-tier system based on feedback about resource exhaustion from uncached clients.
</file>

<file path="issues/html-logger-implementation.md">
# HTML Logger Implementation Plan

## Overview

This document outlines the technical implementation plan for adding an HTML logger
to Langroid's task system. The implementation will create self-contained HTML files
with collapsible log entries, following the specification in `html-logger.md`.

## Architecture

### 1. Core Components

#### 1.1 HTMLLogger Class
Create a new logger class that inherits from or follows the pattern of existing loggers:

```python
class HTMLLogger:
    def __init__(self, filename: str, log_dir: str = "logs"):
        self.file_path = Path(log_dir) / f"{filename}.html"
        self.entries = []
        self._write_header()
    
    def log(self, fields: ChatDocLoggerFields):
        """Add a log entry"""
        entry = self._format_entry(fields)
        self.entries.append(entry)
        self._append_to_file(entry)
    
    def close(self):
        """Finalize the HTML file"""
        self._write_footer()
```

#### 1.2 HTML Template Structure
The HTML file will have three main sections:

1. **Header Section**: Static CSS, JavaScript, and page header
2. **Content Section**: Dynamic log entries
3. **Footer Section**: Closing tags and finalization

### 2. Implementation Steps

#### Step 1: Create HTML Logger Foundation
1. Add `html_logger.py` in `langroid/agent/logging/`
2. Define the `HTMLLogger` class with basic file handling
3. Implement HTML header generation with embedded CSS and JavaScript

#### Step 2: Integrate with Task System
1. Modify `init_loggers` method in `task.py` to include HTML logger option
2. Add configuration flag (e.g., `enable_html_logging` in TaskConfig)
3. Update `log_message` method to send data to HTML logger

#### Step 3: Implement HTML Generation
1. Create entry formatting logic that converts ChatDocLoggerFields to HTML
2. Implement hierarchical structure for collapsible sections
3. Add proper escaping for HTML special characters

#### Step 4: Add JavaScript Functionality
1. Implement toggle functionality for collapsible sections
2. Add "Expand All" / "Collapse All" controls
3. Ensure smooth animations and state management

### 3. Detailed Component Design

#### 3.1 HTML Header Template
```python
HTML_HEADER = """<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>{task_name} - Langroid Task Log</title>
    <style>
        body {{
            background-color: #2b2b2b;
            color: #f0f0f0;
            font-family: 'Consolas', 'Monaco', monospace;
            margin: 0;
            padding: 20px;
        }}
        .header {{
            border: 2px solid #d4a017;
            padding: 10px;
            margin-bottom: 20px;
            color: #d4a017;
        }}
        .entry {{
            margin-bottom: 10px;
            border-left: 3px solid transparent;
        }}
        .entry.user {{ border-left-color: #00ff00; }}
        .entry.assistant {{ border-left-color: #ff6b6b; }}
        .toggle {{
            cursor: pointer;
            user-select: none;
            color: #00ff00;
        }}
        .collapsed .content {{ display: none; }}
        /* More styles... */
    </style>
    <script>
        function toggle(id) {{
            const element = document.getElementById(id);
            element.classList.toggle('collapsed');
            const toggle = element.querySelector('.toggle');
            toggle.textContent = element.classList.contains('collapsed') ? '[+]' : '[-]';
        }}
        /* More JavaScript... */
    </script>
</head>
<body>
    <div class="header">
        <div>{model_info}</div>
        <div>{timestamp} - {message_count} messages</div>
    </div>
    <div id="controls">
        <button onclick="expandAll()">Expand All</button>
        <button onclick="collapseAll()">Collapse All</button>
    </div>
    <div id="content">
"""
```

#### 3.2 Entry Generation Logic
```python
def _format_entry(self, fields: ChatDocLoggerFields) -> str:
    """Convert log fields to HTML entry"""
    entry_id = f"entry_{len(self.entries)}"
    entity_type = fields.responder.upper()
    
    # Build hierarchical structure
    html_parts = [f'<div class="entry {entity_type.lower()}" id="{entry_id}">']
    
    # Add entity header
    if fields.task_name and fields.task_name != "root":
        html_parts.append(f'<div class="entity-header">{fields.task_name} → {entity_type}</div>')
    else:
        html_parts.append(f'<div class="entity-header">{entity_type}</div>')
    
    # Add collapsible sections
    if fields.tool:
        html_parts.append(self._format_tool_section(fields))
    
    # Add main content
    if fields.content:
        html_parts.append(self._format_content_section(fields.content))
    
    html_parts.append('</div>')
    return '\n'.join(html_parts)
```

#### 3.3 Tool Section Formatting
```python
def _format_tool_section(self, fields: ChatDocLoggerFields) -> str:
    """Format tool calls with proper nesting"""
    tool_id = f"tool_{self.tool_counter}"
    self.tool_counter += 1
    
    # Parse tool information
    tool_name = fields.tool
    tool_type = fields.tool_type
    
    # Build tool section HTML
    return f"""
    <div class="tool-section">
        <div class="toggle" onclick="toggle('{tool_id}')">[+]</div>
        <span class="tool-name">{tool_name}({self._format_tool_params(fields)})</span>
        <div id="{tool_id}" class="tool-content collapsed">
            <!-- Tool result and raw call details -->
        </div>
    </div>
    """
```

### 4. Integration Points

#### 4.1 Task Configuration
Add to `TaskConfig`:
```python
class TaskConfig(BaseModel):
    # ... existing fields ...
    enable_html_logging: bool = True
    html_log_dir: str = "logs"
```

#### 4.2 Logger Initialization
Modify `init_loggers` in `task.py`:
```python
def init_loggers(self, tsv_formatter: logging.Formatter | None = None) -> None:
    # ... existing logger setup ...
    
    if self.config.enable_html_logging:
        from langroid.agent.logging.html_logger import HTMLLogger
        self.html_logger = HTMLLogger(
            filename=self.name or "root",
            log_dir=self.config.html_log_dir
        )
```

#### 4.3 Message Logging
Update `log_message` method:
```python
def log_message(self, resp: ChatDocument) -> None:
    # ... existing logging ...
    
    if hasattr(self, 'html_logger') and self.html_logger:
        fields = ChatDocLoggerFields.create(resp, self.id, self.name)
        self.html_logger.log(fields)
```

### 5. Testing Strategy

#### 5.1 Unit Tests
1. Test HTML generation for various message types
2. Test proper escaping of special characters
3. Test file creation and writing
4. Test JavaScript functionality (via parsing)

#### 5.2 Integration Tests
1. Test with simple single-agent tasks
2. Test with multi-agent tasks and sub-tasks
3. Test with various tool types
4. Test with long-running conversations

#### 5.3 Manual Testing
1. Verify visual appearance matches specification
2. Test collapsible functionality in browsers
3. Test performance with large logs
4. Verify accessibility features

### 6. Implementation Timeline

1. **Phase 1**: Core HTML logger class and basic integration (2-3 hours)
2. **Phase 2**: HTML generation with proper styling (2-3 hours)
3. **Phase 3**: JavaScript functionality and interactivity (1-2 hours)
4. **Phase 4**: Testing and refinement (1-2 hours)

### 7. Key Considerations

#### 7.1 Performance
- Stream writes to avoid memory buildup
- Efficient string concatenation
- Minimal JavaScript for responsiveness

#### 7.2 Security
- Proper HTML escaping to prevent XSS
- No external dependencies (self-contained)
- Safe file path handling

#### 7.3 Compatibility
- Test across major browsers
- Ensure proper UTF-8 encoding
- Handle special characters in content

#### 7.4 Edge Cases
- Empty messages
- Very long content
- Special characters in tool names
- Malformed tool responses
- System messages without content

### 8. File Structure

```
langroid/
├── agent/
│   ├── logging/
│   │   ├── __init__.py
│   │   ├── html_logger.py  # New file
│   │   └── ...
│   └── task.py  # Modified
└── tests/
    └── main/
        └── test_html_logger.py  # New test file
```

### 9. Future Extensions

While out of scope for initial implementation, consider:
- Configuration for color themes
- Export functionality
- Search within logs
- Performance optimizations for very large logs
- Real-time streaming updates

### 10. Success Criteria

The implementation will be considered successful when:
1. HTML logs are generated alongside existing logs
2. All log information is preserved and accessible
3. Collapsible sections work smoothly
4. Visual design matches specification
5. No performance impact on task execution
6. Tests pass and edge cases are handled
</file>

<file path="issues/html-logger.md">
# HTML Logger Specification for Langroid Task System

## Overview

This document specifies the requirements for a new HTML logger that will enhance
the current logging system in Langroid's task.py module. The HTML logger will
produce self-contained HTML files with collapsible entries, providing a more
user-friendly way to navigate complex multi-agent conversations.

## Current State

The Langroid task system currently supports two logging formats:
1. **TSV Logger**: Tab-separated values for structured data analysis
2. **Plain Text Logger**: Rich-formatted text logs with color coding

Both loggers capture comprehensive information about agent interactions, including
task names, responders, message content, and tool usage.

## Requirements

### 1. Output Format

- **File Type**: Self-contained HTML file with embedded CSS and JavaScript
- **File Extension**: `.html`
- **File Naming**: Same pattern as existing loggers: `{task_name}.html`
- **Encoding**: UTF-8

### 2. Visual Structure

#### 2.1 Overall Layout
- Dark theme with dark gray/black background (#2b2b2b or similar)
- Monospace font for consistency with terminal output
- Fixed header showing model info and timestamp
- Responsive design that works on various screen sizes
- Golden/amber accent color for headers and borders (#d4a017 or similar)

#### 2.2 Fixed Header Section
- Model name and version (e.g., "claude-opus-4-20250514")
- Timestamp of log generation
- Total message count
- Styled with golden border and text

#### 2.3 Collapsible Entries
Each log entry must be collapsible with:
- **Collapsed State**: Shows only the entity type/role
- **Expanded State**: Shows full message content with sub-sections
- **Toggle Control**: [+] and [-] text indicators in square brackets

#### 2.4 Entry Structure
Each entry consists of:
- **Role Header**: Entity type in colored uppercase (USER, ASSISTANT, SYSTEM, etc.)
- **Collapsible Sections**: Each with [+]/[-] toggle:
  - System Prompt (if applicable)
  - Tools (with count)
  - System Reminder (if applicable)
  - Main content

#### 2.5 Color Scheme
- **USER**: Green text (#00ff00 or similar)
- **ASSISTANT**: Red/orange text (#ff6b6b or similar)
- **SYSTEM**: Gray text
- **Tool calls**: Green indicators for [+] toggles
- **Tool results**: Success (✓) in green, Error (✗) in red
- **Code blocks**: Dark background with syntax highlighting

#### 2.6 Tool Display
When expanded, tool calls should show:
- Tool name and parameters in a code block
- Tool result with success/error indicator
- Raw tool call details (collapsible sub-section)

Example structure:
```
ASSISTANT
[+] System Reminder

I'll read the langroid-llms.txt file to see what it contains.

  [+] Read(/Users/pchalasani/Git/claude-code-play/langroid-llms.txt)
  [+] Tool Result ✓
  [-] Raw Tool Call
      {
        "type": "tool_use",
        "id": "toolu_0184van1ug4T6kAj7a8SkaKp",
        "name": "Read",
        "input": {
          "file_path": "/Users/pchalasani/Git/claude-code-play/langroid-llms.txt"
        }
      }
```

### 3. Functionality

#### 3.1 User Controls
- **Expand/Collapse Individual**: Click on entry header or toggle button
- **Expand All**: Button to expand all entries
- **Collapse All**: Button to collapse all entries
- **Search**: Basic text search functionality (optional enhancement)

#### 3.2 State Persistence
- Collapse/expand state should be maintained during the session
- No requirement for persistence across page reloads

### 4. Data Representation

The HTML logger should capture all information currently logged by the plain
text logger and organize it hierarchically:

#### 4.1 Primary Level (Always Visible)
- Entity/Role name (USER, ASSISTANT, AGENT, etc.)
- Task name prefix if not "root"

#### 4.2 Collapsible Sections
Each entry may have multiple collapsible sub-sections:
- **System Messages**: System prompts, reminders, etc.
- **Tool Information**: 
  - Tool count in header (e.g., "Tools (17)")
  - Individual tool calls with name and parameters
  - Tool results with success/error indicators
  - Raw tool call JSON (nested collapsible)
- **Message Content**: The actual text content
- **Metadata** (when relevant):
  - Recipient information
  - Blocked entities
  - Mark indicator for final results

#### 4.3 Mapping from Current Log Fields
- `responder` → Entity type (USER, ASSISTANT, etc.)
- `task_name` → Prefix before entity if not "root"
- `sender` + `sender_name` → Combined in display
- `tool_type` + `tool` → Tool section with appropriate formatting
- `content` → Main message content
- `mark` → Special indicator for final results
- `recipient` → Shown in metadata when present
- `block` → Shown in metadata when present

### 5. Integration Requirements

#### 5.1 Implementation Location
- Add to the existing `init_loggers` method in task.py
- Follow the same pattern as TSV and plain text loggers
- Use the same log directory and naming conventions

#### 5.2 Configuration
- HTML logger should be optional
- Controlled via configuration flag or environment variable
- Should not interfere with existing loggers

#### 5.3 Compatibility
- Must work with the existing `log_message` method
- Support the same ChatDocLoggerFields structure
- Handle sub-tasks transparently (no special handling needed)

### 6. Performance Considerations

- Efficient for files with thousands of log entries
- Minimal JavaScript for toggle functionality
- CSS animations should be optional or lightweight
- File size should remain reasonable for large conversations

### 7. Accessibility

- Keyboard navigation support for expanding/collapsing entries
- Clear visual indicators for interactive elements
- Sufficient color contrast for readability
- Screen reader compatible structure

## Example Visual Mock-up

```
claude-opus-4-20250514
7/8/2025, 12:00:50 PM   8 messages
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[+] System Prompt
[+] Tools (17)

USER
[+] System Reminder

Use the read tool to read the langroid-llms.txt and see what it is about

[+] System Reminder

ASSISTANT

I'll read the langroid-llms.txt file to see what it contains.

  [+] Read(/Users/pchalasani/Git/claude-code-play/langroid-llms.txt)
  [+] Tool Result ✗
  [-] Raw Tool Call
      {
        "type": "tool_use",
        "id": "toolu_0184van1ug4T6kAj7a8SkaKp",
        "name": "Read",
        "input": {
          "file_path": "/Users/pchalasani/Git/claude-code-play/langroid-llms.txt"
        }
      }

ASSISTANT

The file is quite large (3.3MB). Let me read it in smaller chunks to understand its content.

  [-] Read(/Users/pchalasani/Git/claude-code-play/langroid-llms.txt)
      {
        "file_path": "/Users/pchalasani/Git/claude-code-play/langroid-llms.txt",
        "limit": 100
      }
  [+] Tool Result ✓
  [+] Raw Tool Call
```

## Future Enhancements (Out of Scope)

These features are not required for the initial implementation but could be
added later:
- Filtering by entity type, task name, or tool
- Export to other formats
- Real-time log streaming
- Syntax highlighting for code in messages
- Timestamp display options
- Log entry grouping by conversation threads
</file>

<file path="issues/issue-919-llamacpp-embeddings.md">
# Issue #919: llama.cpp Embeddings Support

## Background

User reported issues using llama.cpp server for local embeddings with Langroid. The error occurred when using `LlamaCppServerEmbeddingsConfig`:

```
TypeError: list indices must be integers or slices, not str
```

This happened at line 466 in `langroid/embedding_models/models.py`:
```python
embeddings = response.json()["embedding"]
```

## Investigation Summary

### Can llama.cpp Generate Embeddings?

**YES!** llama.cpp supports embeddings in two ways:

1. **Dedicated embedding models** (RECOMMENDED):
   - nomic-embed-text-v1.5 (768 dims)
   - nomic-embed-text-v2-moe
   - nomic-embed-code
   - Other GGUF embedding models

2. **Regular LLMs** (works but not optimal):
   - gpt-oss-20b, gpt-oss-120b
   - Llama models
   - By extracting internal representations

### How to Enable

Start llama-server with the `--embeddings` flag:

```bash
./llama-server -ngl 100 -c 2048 \
  -m ~/nomic-embed-text-v1.5.Q8_0.gguf \
  --host localhost --port 8080 \
  --embeddings -b 2048 -ub 2048
```

## llama.cpp Embedding Endpoints

llama.cpp provides multiple embedding endpoints with different response formats:

### 1. Native `/embedding` endpoint

**Request:**
```json
{
  "content": "your text here"
}
```

**Response:**
```json
{
  "embedding": [0.1, 0.2, 0.3, ...]
}
```

### 2. OpenAI-compatible `/v1/embeddings` endpoint

**Request:**
```json
{
  "input": "your text here",
  "model": "model-name"
}
```

**Response:**
```json
{
  "object": "list",
  "data": [
    {
      "object": "embedding",
      "embedding": [0.1, 0.2, 0.3, ...],
      "index": 0
    }
  ],
  "model": "model-name",
  "usage": {
    "prompt_tokens": 5,
    "total_tokens": 5
  }
}
```

## The Problem

The original Langroid code expected only format #1 (native):
```python
embeddings = response.json()["embedding"]
```

However, llama.cpp can return **different formats** depending on:
- Endpoint used (`/embedding` vs `/v1/embeddings`)
- Server version/configuration
- Batch mode settings

The error indicated that `response.json()` returned a **list**, not a **dict**, suggesting llama.cpp returned an array format.

## Discovered Response Formats

Through investigation, we identified **5 possible response formats**:

1. **Native format**: `{"embedding": [floats]}`
2. **Array format**: `[{"embedding": [floats]}]`
3. **Double-nested**: `[{"embedding": [[floats]]}]`
4. **OpenAI-compatible**: `{"data": [{"embedding": [floats]}]}`
5. **Dict-nested**: `{"embedding": [[floats]]}`

## Our Solution

### Implementation

Added a robust `_extract_embedding()` method in `langroid/embedding_models/models.py` (lines 483-544) that:

1. Tries each format in order
2. Validates the extracted embedding is a list of floats
3. Provides clear error messages if format is unrecognized

```python
def _extract_embedding(
    self, response_json: dict[str, Any] | list[Any]
) -> List[int | float]:
    """
    Extract embedding vector from llama.cpp response.

    Handles multiple response formats:
    1. Native /embedding: {"embedding": [floats]}
    2. Array format: [{"embedding": [floats]}]
    3. Double-nested: [{"embedding": [[floats]]}]
    4. OpenAI /v1/embeddings: {"data": [{"embedding": [floats]}]}
    5. Nested in dict: {"embedding": [[floats]]}

    Args:
        response_json: The JSON response from llama.cpp server

    Returns:
        List of floats representing the embedding vector

    Raises:
        ValueError: If response format is not recognized
    """
    # Implementation handles all 5 formats...
```

### Modified `generate_embedding()` method

Changed from:
```python
embeddings = response.json()["embedding"]
```

To:
```python
embeddings = self._extract_embedding(response.json())
```

## Testing

Created comprehensive unit tests in `tests/extras/test_llamacpp_embedding_formats.py`:

- ✅ test_native_format
- ✅ test_array_format
- ✅ test_double_nested_array_format
- ✅ test_openai_compatible_format
- ✅ test_nested_in_dict_format
- ✅ test_invalid_format_raises_error
- ✅ test_generate_embedding_with_native_format (mocked)
- ✅ test_generate_embedding_with_array_format (mocked)
- ✅ test_generate_embedding_with_openai_format (mocked)
- ✅ test_generate_embedding_http_error

**All tests pass** ✅
**Linting and type checking pass** ✅

## Comparison with PR #920

### PR #920 Approach

Changed:
```python
embeddings = response.json()["embedding"]
```

To:
```python
embeddings = response.json()[0]["embedding"][0]
```

### Issues with PR #920

1. **Too specific**: Only handles ONE format: `[{"embedding": [[floats]]}]`
2. **Logic error**: The double `[0]` indexing would extract a single float, not the full embedding vector
3. **Would fail validation**: The existing validation expects a list of floats
4. **No tests**: No unit tests provided
5. **No documentation**: No explanation of what format is expected

### Our Solution Advantages

1. **Handles 5 different formats** automatically
2. **Backwards compatible**: Works with existing deployments
3. **Well-tested**: 10 unit tests covering all scenarios
4. **Well-documented**: Clear docstring explaining all formats
5. **Robust error messages**: Helps users debug configuration issues

## Usage Example

### Configuration

```python
from langroid.embedding_models.models import LlamaCppServerEmbeddingsConfig
from langroid.vector_store.qdrantdb import QdrantDBConfig

embed_cfg = LlamaCppServerEmbeddingsConfig(
    api_base="http://localhost:8080",  # Your llama.cpp server
    dims=768,  # Match your embedding model dimensions
    context_length=2048,
    batch_size=2048,
)

vecdb_config = QdrantDBConfig(
    collection_name="my-docs",
    embedding=embed_cfg,
    storage_path=".qdrant/",
)
```

### Running llama-server

```bash
# For dedicated embedding model (RECOMMENDED)
./llama-server -ngl 100 -c 2048 \
  -m ~/nomic-embed-text-v1.5.Q8_0.gguf \
  --embeddings -b 2048 -ub 2048 \
  --host localhost --port 8080

# For LLM-based embeddings (gpt-oss example)
./llama-server -ngl 99 \
  -m ~/.cache/llama.cpp/gpt-oss-20b.gguf \
  --embeddings \
  --host localhost --port 8080
```

## Recommendations

### For Users

1. **Use dedicated embedding models** like nomic-embed-text-v1.5 for best results
2. **Match dimensions** in config to your embedding model
3. **Use the `--embeddings` flag** when starting llama-server
4. **Check server logs** if you encounter issues

### For Langroid

1. ✅ **Implemented**: Robust format detection in `_extract_embedding()`
2. ✅ **Tested**: Comprehensive unit tests
3. ✅ **Documented**: Clear docstrings and examples
4. **Consider**: Adding example in `examples/docqa/` using local embeddings
5. **Consider**: Adding to documentation/tutorials

## Files Modified

- `langroid/embedding_models/models.py` - Added `_extract_embedding()` method
- `tests/extras/test_llamacpp_embedding_formats.py` - New comprehensive test suite

## References

- Issue #919: https://github.com/langroid/langroid/issues/919
- PR #920: https://github.com/langroid/langroid/pull/920
- llama.cpp discussion #7712: https://github.com/ggml-org/llama.cpp/discussions/7712
- nomic-embed models: https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF
- Langroid docs: `docs/notes/llama-cpp-embeddings.md`

## Conclusion

**Issue #919 is now resolved** with a robust, well-tested solution that handles all known llama.cpp embedding response formats. Users can now use local embeddings with llama.cpp without worrying about response format variations.

**PR #920 is not needed** as our solution is more comprehensive and handles all cases, not just one specific format.
</file>

<file path="issues/llm-client-caching-phase1-summary.md">
# Phase 1 Implementation Summary: Client Caching

## Changes Made

### 1. Created `langroid/language_models/client_cache.py`

A new module implementing singleton pattern for LLM clients with the following features:

- **Consistent with existing caching**: Uses SHA256 hashing for cache keys, matching the approach in `OpenAIGPT._cache_lookup`
- **Wrapper functions** for each client type:
  - `get_openai_client()` / `get_async_openai_client()`
  - `get_groq_client()` / `get_async_groq_client()`
  - `get_cerebras_client()` / `get_async_cerebras_client()`
- **Configuration-based caching**: Clients are cached based on their full configuration (API key, base URL, timeout, headers, etc.)
- **Lifecycle management**: Uses `atexit` hook for cleanup and weak references to track clients

### 2. Key Implementation Details

#### Cache Key Generation
```python
def _get_cache_key(client_type: str, **kwargs: Any) -> str:
    # Convert kwargs to sorted string representation
    sorted_kwargs_str = str(sorted(kwargs.items()))
    
    # Create raw key combining client type and sorted kwargs
    raw_key = f"{client_type}:{sorted_kwargs_str}"
    
    # Hash the key for consistent length and to handle complex objects
    hashed_key = hashlib.sha256(raw_key.encode()).hexdigest()
    
    return hashed_key
```

This approach:
- Ensures deterministic keys through sorting
- Handles complex objects via string representation
- Produces fixed-length keys (64 chars)
- Matches the existing Redis cache key generation pattern

### 3. Created Comprehensive Tests

`tests/main/test_client_cache.py` includes tests for:
- Singleton behavior (same config returns same client)
- Different configurations return different clients
- Different client types are cached separately
- Proper handling of timeout objects and headers
- Type differences are preserved (e.g., `30` vs `30.0` are different)

### 4. All Quality Checks Pass
- ✅ All 9 tests pass
- ✅ Type checking passes (mypy)
- ✅ Linting passes (ruff, black)

## Design Decisions

1. **Used SHA256 hashing instead of tuple keys**: More robust for complex objects and consistent with existing caching approach
2. **Type strictness**: `30` and `30.0` create different cache entries - better to be overly strict than risk bugs
3. **Weak references**: Allow garbage collection of unused clients while maintaining cleanup capability
4. **Simple atexit cleanup**: Accepted that async clients will be cleaned by OS on exit

## Next Steps (Phase 2)

Update `OpenAIGPT.__init__` to use these wrapper functions instead of directly creating clients:
```python
# Current
self.client = OpenAI(api_key=self.api_key, ...)

# New  
from langroid.language_models.client_cache import get_openai_client
self.client = get_openai_client(api_key=self.api_key, ...)
```

This will require careful updating of all client creation locations in `openai_gpt.py`.
</file>

<file path="issues/llm-client-caching-phase2-summary.md">
# Phase 2 Implementation Summary: OpenAIGPT Integration

## Changes Made

### 1. Updated `langroid/language_models/openai_gpt.py`

- **Added imports** for client cache wrapper functions
- **Replaced direct client instantiation** with wrapper functions:
  - `Groq()` → `get_groq_client()`
  - `AsyncGroq()` → `get_async_groq_client()`
  - `Cerebras()` → `get_cerebras_client()`
  - `AsyncCerebras()` → `get_async_cerebras_client()`
  - `OpenAI()` → `get_openai_client()`
  - `AsyncOpenAI()` → `get_async_openai_client()`

### 2. Fixed Async Client Cleanup

Updated `_cleanup_clients()` to properly handle async clients by checking if `close()` is a coroutine function and skipping await (since atexit can't handle async).

### 3. Created Integration Tests

`tests/main/test_openai_gpt_client_cache.py` with tests verifying:
- Multiple OpenAIGPT instances with same config share clients
- Different configurations create different clients
- Works correctly for OpenAI, Groq, and Cerebras models
- Different base URLs and headers create different clients

## Results

### Before (Anti-pattern)
```python
# Creating 100 agents = 100 OpenAI clients
for row in data[:100]:
    agent = ChatAgent(config)  # Each creates new OpenAI client
    result = agent.run(row)
```

### After (With caching)
```python
# Creating 100 agents = 1 OpenAI client (reused)
for row in data[:100]:
    agent = ChatAgent(config)  # Reuses existing OpenAI client
    result = agent.run(row)
```

## Testing Results

- ✅ All 9 client cache unit tests pass
- ✅ All 6 OpenAIGPT integration tests pass
- ✅ Existing LLM tests continue to pass
- ✅ Type checking passes
- ✅ Linting passes

## Benefits

1. **Resource Efficiency**: Dramatically reduces file descriptor usage
2. **Performance**: Eliminates repeated client initialization overhead
3. **Transparent**: No API changes required - existing code benefits automatically
4. **Configurable**: Each unique configuration gets its own cached client
5. **Safe**: Thread-safe implementation with proper cleanup

## Implementation Notes

- Used SHA256 hashing for cache keys (consistent with existing Redis cache)
- Handles all configuration parameters (API key, base URL, timeout, headers, etc.)
- Async client cleanup deferred to OS (atexit can't await)
- Weak references allow garbage collection when clients no longer needed
</file>

<file path="issues/llm-client-caching-test-summary.md">
# Client Caching Test Summary

## Tests Created

### 1. Unit Tests (`test_client_cache.py`)
- **Purpose**: Test the basic caching functionality
- **Coverage**: 
  - Singleton behavior for same configuration
  - Different clients for different configurations
  - Proper handling of all client types (OpenAI, Groq, Cerebras)
  - Cache key generation with complex types

### 2. Integration Tests (`test_openai_gpt_client_cache.py`)
- **Purpose**: Test OpenAIGPT integration with caching
- **Coverage**:
  - Multiple OpenAIGPT instances share clients
  - Different configs create different clients
  - Works for all model types (OpenAI, Groq, Cerebras)

### 3. Stress Tests (`test_client_cache_stress.py`)
- **Purpose**: Demonstrate resource usage improvements
- **Tests**:
  - `test_many_agents_with_caching`: Shows 100 agents share 1 client
  - `test_many_agents_different_configs`: Shows proper separation by config
  - `test_memory_efficiency`: Demonstrates memory savings
  - `test_client_instance_comparison`: Direct comparison with/without caching

### 4. Demonstration Test (`test_client_cache_demo.py`)
- **Purpose**: Clear demonstration of the fix for the exact user scenario
- **Key Results**:

#### With Client Caching:
- 100 ChatAgent instances → 1 shared client pair
- File descriptors saved: ~297
- Memory saved: ~148.5 MB
- Creation time: 0.60 seconds

#### Without Client Caching (simulated):
- 100 ChatAgent instances → 100 client pairs
- File descriptors used: ~300
- Extra memory used: ~148.5 MB
- Risk of "Too many open files" errors

## Test Results Summary

All tests demonstrate that the client caching implementation:

1. **Prevents resource exhaustion**: 100 agents use 1 client instead of 100
2. **Maintains correctness**: Different configurations still get different clients
3. **Is transparent**: No API changes needed
4. **Provides significant savings**:
   - 50x reduction in client instances
   - ~297 file descriptors saved for 100 agents
   - ~148.5 MB memory saved for 100 agents

The stress tests confirm that the implementation successfully addresses the "too many open files" issue that was occurring when creating many agents in a loop.
</file>

<file path="issues/llm-client-caching.md">
# LLM Client Connection Pool Exhaustion Issue

## Problem Statement

When using Langroid in multi-agent systems where agents are created dynamically (e.g., one agent per data row), each agent creates its own LLM client instance (OpenAI, Groq, or Cerebras). This pattern leads to connection pool exhaustion, resulting in "too many open files" errors and degraded performance.

## Current Behavior

### Client Creation Flow
1. Each `ChatAgent` instantiates its own `OpenAIGPT` instance
2. Each `OpenAIGPT` instance creates new client objects:
   - For Groq models: Creates `Groq()` and `AsyncGroq()` clients
   - For Cerebras models: Creates `Cerebras()` and `AsyncCerebras()` clients  
   - For OpenAI/others: Creates `OpenAI()` and `AsyncOpenAI()` clients
3. These clients maintain their own connection pools via httpx

### Problem Scenario
```python
# Anti-pattern: Creating many agents
for row in data[:100]:  # 100 rows
    agent = ChatAgent(config)  # Creates new OpenAI client
    result = agent.run(row)    # Makes API calls
    # Agent goes out of scope but connections may linger
```

This creates 100 separate OpenAI clients, each with its own connection pool.

## Impact

1. **Resource Exhaustion**: Each client maintains open connections, leading to file descriptor limits
2. **Performance Degradation**: Connection establishment overhead for each new client
3. **Potential API Rate Limiting**: Multiple clients may trigger more aggressive rate limiting
4. **Memory Usage**: Each client instance consumes memory for connection pools

## Root Cause

The issue stems from:
1. Lack of client reuse across agent instances
2. Connection pools not being properly closed when agents are garbage collected
3. The anti-pattern of creating many short-lived agents instead of reusing agents

## Constraints

1. **API Compatibility**: Solution must not break existing Langroid API
2. **Configuration Flexibility**: Different agents may need different configurations (API keys, base URLs, timeouts)
3. **Thread Safety**: Clients must be safely shareable across multiple agents
4. **Async Support**: Must handle both sync and async client variants

## Critical Considerations

### 1. Configuration Variations
Different agents in the same system might require different client configurations:
- **Different API Keys**: Agent A might use one OpenAI key, Agent B another
- **Different Base URLs**: Some agents might use standard OpenAI, others might use Azure OpenAI
- **Different Timeouts**: Long-running tasks might need higher timeouts
- **Different Headers**: Custom headers for different use cases

**Implication**: We cannot have just one singleton per client type. We need to cache clients based on their full configuration, creating a new client only when a unique configuration is encountered.

### 2. Thread Safety
Multiple agents might run concurrently and share the same client instance:
- The httpx library (used by OpenAI, Groq, Cerebras clients) is designed to be thread-safe
- Connection pools in httpx can handle concurrent requests
- No additional locking should be needed for client access

**Implication**: Shared clients can be used safely across multiple threads/agents without synchronization overhead.

### 3. Lifecycle Management
Proper cleanup of singleton clients is crucial:
- **When to close**: Clients hold network resources that should be released
- **Garbage collection**: Need to ensure clients can be GC'd when no longer needed
- **Application shutdown**: Should close all clients gracefully on exit

**Implications**: 
- Consider using weak references to allow garbage collection of unused clients
- Implement `atexit` hooks for graceful shutdown
- May need a manual cleanup mechanism for long-running applications
- Monitor for memory leaks from accumulating cached clients with unique configs

## Proposed Solution: Client Singleton Pattern

### Approach
Implement a caching layer that returns singleton clients based on configuration:

1. **Wrapper Functions**: Replace direct client instantiation with wrapper functions:
   - `get_openai_client(config) -> OpenAI`
   - `get_groq_client(config) -> Groq`
   - `get_cerebras_client(config) -> Cerebras`
   - Similar for async variants

2. **Configuration-Based Caching**: Cache clients keyed by their configuration parameters:
   - API key
   - Base URL
   - Timeout
   - Headers
   - Organization (for OpenAI)

3. **Implementation Location**: In `langroid/language_models/openai_gpt.py`, replace:
   ```python
   # Current
   self.client = OpenAI(api_key=self.api_key, ...)
   
   # Proposed
   self.client = get_openai_client(api_key=self.api_key, ...)
   ```

### Benefits
- Reduces client instances from N (number of agents) to M (unique configurations)
- No API changes required
- Follows OpenAI best practices for client reuse
- Transparent to existing code

### Alternative Solutions Considered

1. **Agent Pooling**: Reuse agents instead of creating new ones
   - Pros: Most efficient
   - Cons: Requires significant API changes

2. **Explicit Client Registry**: Pass shared clients to agents
   - Pros: Explicit control
   - Cons: Breaks existing API, requires user awareness

3. **Connection Limit Configuration**: Reduce connection pool sizes
   - Pros: Simple
   - Cons: Doesn't address root cause, may hurt performance

## Success Criteria

1. Creating 100+ agents should not cause "too many open files" errors
2. Memory usage should remain stable with many agents
3. No breaking changes to existing Langroid API
4. Performance improvement for multi-agent scenarios

## Implementation Notes

- httpx clients (used by OpenAI/Groq/Cerebras) are thread-safe
- Consider using weak references to allow garbage collection
- May need cleanup hooks (atexit) for proper shutdown
- Should add logging for cache hits/misses for debugging

## References

- OpenAI Cookbook: Best practices recommend reusing client instances
- httpx documentation: Connection pooling behavior
- Python file descriptor limits and ulimit settings
</file>

<file path="issues/pr-882-cached-tokens-improvements.md">
# PR #882: Cached Tokens Support - Improvements

## Summary
Enhanced PR #882 which adds cached token tracking to LLMTokenUsage. Made several improvements including cleanup of unused code, bug fixes, added tests, and new model support.

## Changes

### 1. Code Cleanup
- Removed unused `chat_cost_per_1k_tokens` and `completion_cost_per_1k_tokens` fields from `LLMConfig` in `base.py`
- These fields were superseded by the ModelInfo approach but were still being updated unnecessarily

### 2. Bug Fixes
- Fixed type error in `openai_gpt.py` when extracting `prompt_tokens_details` from API responses
- Added proper type annotation and type checking to handle cases where the field might not be a dict

### 3. Added Tests
- `test_cached_tokens_tracking()`: Verifies cached tokens are properly tracked in API responses and cost calculations work correctly
- `test_cached_tokens_in_llm_response()`: Tests the LLMTokenUsage class directly including string representation and reset functionality

### 4. Added Gemini 2.5 Model Support
- Fixed `GEMINI_2_5_PRO` enum to map to `"gemini-2.5-pro"` instead of experimental version
- Added new enums: `GEMINI_2_5_FLASH` and `GEMINI_2_5_FLASH_LITE_PREVIEW`
- Added complete ModelInfo entries with proper costs and parameters:
  - **Gemini 2.5 Pro**: 1M context, $1.25/$0.31/$10.00 per million tokens
  - **Gemini 2.5 Flash**: 1M context, $0.30/$0.075/$2.50 per million tokens  
  - **Gemini 2.5 Flash Lite Preview**: 64K context, $0.10/$0.025/$0.40 per million tokens

## Testing
- All existing tests pass
- New tests verify cached token functionality
- Code passes all linting and type checking
</file>

<file path="issues/pr-openai-client-caching.md">
# OpenAI Client Connection Management

## Problem
Creating many agents (e.g., 100 agents for 100 data rows) leads to "too many open files" errors due to each agent creating its own HTTP client, exhausting file descriptors.

## Solution
Implemented client caching/singleton pattern to reuse HTTP clients across multiple agent instances with the same configuration.

## Changes

### 1. Client Caching Module (`langroid/language_models/client_cache.py`)
- Singleton pattern for HTTP client reuse
- SHA256-based cache keys for configuration
- Wrapper functions for each client type (OpenAI, Groq, Cerebras)
- Lifecycle management with `atexit` hooks

### 2. OpenAIGPT Integration
- Added `use_cached_client: bool = True` config parameter
- Updated client creation to use wrapper functions when caching enabled
- Allows disabling for testing/special cases

### 3. ChatAgent Cleanup
- Updated `__del__` method to avoid closing shared clients
- Clients now managed centrally via client_cache module

### 4. Comprehensive Tests
- Tests for singleton behavior across all client types
- Verification of concurrent async usage
- Tests for model prefix routing (groq/, cerebras/, etc.)
- Regression tests with `use_cached_client` flag

## Benefits
- Prevents resource exhaustion when creating many agents
- Improves performance through connection pooling
- Backward compatible with opt-out capability
- Thread-safe for concurrent usage
</file>

<file path="issues/pr-qdrant-lock-fix.md">
# Fix QdrantDB Lock File Issue

## Problem
When using QdrantDB with local storage, file lock conflicts occurred when:
1. A QdrantDB instance was created but not properly closed
2. Another part of the code tried to create a new QdrantDB instance with the same storage path
3. Qdrant would detect the `.lock` file and create a new storage directory (e.g., `./qdrant_data.new`)

## Solution
1. **Added `close()` method to QdrantDB** - Calls the underlying client's close method to release the file lock
2. **Added context manager support** - Implemented `__enter__` and `__exit__` for automatic cleanup
3. **Fixed DocChatAgent's `clear()` method** - Now closes the old vecdb before creating a new one

## Usage
```python
# Option 1: Explicit close
vecdb = QdrantDB(config)
vecdb.clear_all_collections(really=True)
vecdb.close()  # Release the lock

# Option 2: Context manager (automatic cleanup)
with QdrantDB(config) as vecdb:
    vecdb.clear_all_collections(really=True)
    # Automatically closed when exiting context
```

## Changes
- `langroid/vector_store/qdrantdb.py`: Added `close()`, `__enter__`, `__exit__` methods
- `langroid/agent/special/doc_chat_agent.py`: Fixed `clear()` to close old vecdb instance

This fix prevents the proliferation of `.new` directories when using QdrantDB with local storage.
</file>

<file path="issues/qdrant-lock-issue-spec-changes.md">
# QdrantDB Lock File Conflict Issue - Changes and Best Practices

## Summary of Changes

This document describes the changes made to resolve the QdrantDB lock file conflict issue described in `qdrant-lock-issue-spec.md`.

## Problem Recap

When using QdrantDB with local storage, a file lock conflict occurred when:
1. A QdrantDB instance was created (e.g., to clear collections)
2. The instance was not properly disposed/closed
3. Another part of the code tried to create a new QdrantDB instance
4. Qdrant detected the `.lock` file and created a new storage directory (e.g., `./qdrant_data.new`)

## Implemented Solution

### 1. Added `close()` Method

Added an explicit `close()` method to the QdrantDB class:

```python
def close(self) -> None:
    """
    Close the QdrantDB client and release any resources (e.g., file locks).
    This is especially important for local storage to release the .lock file.
    """
    if hasattr(self.client, "close"):
        # QdrantLocal has a close method that releases the lock
        self.client.close()
        logger.info(f"Closed QdrantDB connection for {self.config.storage_path}")
```

### 2. Added Context Manager Support

Implemented `__enter__` and `__exit__` methods to support Python's context manager protocol:

```python
def __enter__(self) -> "QdrantDB":
    """Context manager entry."""
    return self

def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
    """Context manager exit - ensure cleanup even if an exception occurred."""
    self.close()
```

### 3. Added Type Import

Added `Any` to the type imports to support the context manager type hints:

```python
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, TypeVar
```

## Best Practices for Using QdrantDB

### Important Note

The underlying `qdrant_client` library does not implement the context manager protocol. However, **Langroid's QdrantDB wrapper now provides context manager support** to ensure proper cleanup of resources, especially the file lock used by QdrantLocal.

### Recommended: Use Context Manager (Most Pythonic)

The context manager approach is the **recommended best practice** for Langroid's QdrantDB as it guarantees cleanup even if exceptions occur:

```python
from langroid.vector_store.qdrantdb import QdrantDB, QdrantDBConfig

config = QdrantDBConfig(
    cloud=False,
    collection_name="my_collection",
    storage_path="./qdrant_data",
)

# Recommended approach
with QdrantDB(config) as vecdb:
    # Use the vector database
    vecdb.add_documents(documents)
    results = vecdb.similar_texts_with_scores("query text", k=5)
    vecdb.clear_empty_collections()
    # Automatically closed when exiting the context
```

### Alternative: Explicit `close()` Method

If you cannot use a context manager (e.g., when the QdrantDB instance needs to persist across multiple methods), use explicit `close()`:

```python
class MyDocProcessor:
    def __init__(self):
        config = QdrantDBConfig(
            cloud=False,
            collection_name="my_collection",
            storage_path="./qdrant_data",
        )
        self.vecdb = QdrantDB(config)
    
    def process_documents(self, docs):
        self.vecdb.add_documents(docs)
    
    def search(self, query):
        return self.vecdb.similar_texts_with_scores(query, k=5)
    
    def cleanup(self):
        # Important: Call this when done
        self.vecdb.close()
```

### When Using with DocChatAgent

When using QdrantDB with DocChatAgent, the agent manages the vector store lifecycle, so you don't need to worry about closing it manually:

```python
from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig

# DocChatAgent manages the QdrantDB lifecycle
agent = DocChatAgent(
    DocChatAgentConfig(
        vecdb=QdrantDBConfig(
            cloud=False,
            collection_name="doc_chat",
            storage_path="./qdrant_data",
        )
    )
)
# The agent will handle cleanup appropriately
```

### For Temporary Operations

For one-off operations like clearing collections, always use context manager:

```python
# Clear all collections
with QdrantDB(config) as vecdb:
    vecdb.clear_all_collections(really=True, prefix="temp_")

# Clear and recreate
with QdrantDB(config) as vecdb:
    vecdb.delete_collection("old_collection")
    vecdb.create_collection("new_collection", replace=True)
```

## Important Notes

1. **Cloud Storage**: This issue only affects local storage (`cloud=False`). When using Qdrant cloud service, file locking is not used.

2. **Backward Compatibility**: Existing code will continue to work without changes, but may show warnings about lock conflicts and create `.new` directories.

3. **Multiple Processes**: If you genuinely need multiple processes to access the same Qdrant storage simultaneously, use Qdrant server instead of local storage.

## Testing

Comprehensive tests were added to verify the fix:
- `tests/main/test_qdrant_lock_release.py` - Unit tests for close() and context manager
- `tests/main/test_qdrant_lock_scenario.py` - Reproduces the exact issue scenario
- `tests/main/test_qdrant_warning_capture.py` - Captures and verifies warning messages

All tests pass and confirm that:
- Without proper cleanup: `.new` directories are created (the bug)
- With `close()` or context manager: No `.new` directories (fixed)

## Migration Guide

If you have existing code that creates temporary QdrantDB instances:

**Before (problematic):**
```python
vecdb = QdrantDB(config)
vecdb.clear_all_collections(really=True)
# Lock file remains, causing issues
```

**After (fixed):**
```python
# Option 1: Context manager (preferred)
with QdrantDB(config) as vecdb:
    vecdb.clear_all_collections(really=True)

# Option 2: Explicit close
vecdb = QdrantDB(config)
vecdb.clear_all_collections(really=True)
vecdb.close()
```

## Why This Matters

While the `qdrant_client` library handles some cleanup via its `__del__` method, this is not reliable because:
1. Python's garbage collector doesn't guarantee when `__del__` will be called
2. In some cases (circular references, interpreter shutdown), `__del__` may not be called at all
3. The file lock remains until the process ends, preventing other instances from using the same storage

By adding explicit `close()` and context manager support to Langroid's QdrantDB wrapper, we ensure:
- Immediate release of the file lock when done
- No proliferation of `.new` directories
- Predictable resource cleanup
- Better development experience (no need to manually delete lock files)

## Conclusion

The QdrantDB lock file issue has been resolved by adding proper resource cleanup mechanisms to Langroid's QdrantDB wrapper. While the underlying `qdrant_client` doesn't provide context manager support, Langroid now offers both context manager and explicit `close()` methods. The context manager approach is the recommended best practice as it ensures cleanup even in error scenarios. For cases where context managers aren't suitable, the explicit `close()` method provides a reliable alternative.
</file>

<file path="langroid/agent/callbacks/__init__.py">

</file>

<file path="langroid/agent/special/arangodb/__init__.py">

</file>

<file path="langroid/agent/special/arangodb/system_messages.py">
from langroid.agent.special.arangodb.tools import (
    aql_creation_tool_name,
    aql_retrieval_tool_name,
    arango_schema_tool_name,
)
from langroid.agent.tools.orchestration import DoneTool

done_tool_name = DoneTool.default_value("request")

arango_schema_tool_description = f"""
`{arango_schema_tool_name}` tool/function-call to find the schema
of the graph database, or for some SPECIFIC collections, i.e. get information on 
(document and edge), their attributes, and graph definitions available in your
ArangoDB database. You MUST use this tool BEFORE attempting to use the
`{aql_retrieval_tool_name}` tool/function-call, to ensure that you are using the
correct collection names and attributes in your `{aql_retrieval_tool_name}` tool.
"""

aql_retrieval_tool_description = f"""
`{aql_retrieval_tool_name}` tool/function-call to retrieve information from 
  the database using AQL (ArangoDB Query Language) queries, to answer
  the user's questions, OR for you to learn more about the SCHEMA of the database.
"""

aql_creation_tool_description = f"""
`{aql_creation_tool_name}` tool/function-call to execute AQL query that creates
documents/edges in the database.
"""

aql_retrieval_query_example = """
EXAMPLE:
Suppose you are asked this question "Does Bob have a father?".
Then you will go through the following steps, where YOU indicates
the message YOU will be sending, and RESULTS indicates the RESULTS
you will receive from the helper executing the query:

1. YOU:
    {{ "request": "aql_retrieval_tool",
      "aql_query": "FOR v, e, p in ... [query truncated for brevity]..."}}

    2. RESULTS:
    [.. results from the query...]
    3. YOU: [ since results were not satisfactory, you try ANOTHER query]
    {{ "request": "aql_retrieval_tool",
    "aql_query": "blah blah ... [query truncated for brevity]..."}}
    }}
    4. RESULTS:
    [.. results from the query...]
    5. YOU: [ now you have the answer, you can generate your response ]
    The answer is YES, Bob has a father, and his name is John.
"""

aql_query_instructions = """
When writing AQL queries:
1. Use the exact property names shown in the schema
2. Pay attention to the 'type' field of each node
3. Note that all names are case-sensitive:
   - collection names
   - property names
   - node type values
   - relationship type values
4. Always include type filters in your queries, e.g.:
   FILTER doc.type == '<type-from-schema>'

The schema shows:
- Collections (usually 'nodes' and 'edges')
- Node types in each collection
- Available properties for each node type
- Relationship types and their properties

Examine the schema carefully before writing queries to ensure:
- Correct property names
- Correct node types
- Correct relationship types

You must be smart about using the right collection names and attributes
based on the English description. If you are thinking of using a collection
or attribute that does not exist, you are probably on the wrong track,
so you should try your best to answer based on existing collections and attributes.
DO NOT assume any collections or graphs other than those above.
"""

tool_result_instruction = """
REMEMBER:
[1]  DO NOT FORGET TO USE ONE OF THE AVAILABLE TOOLS TO ANSWER THE USER'S QUERY!!
[2] When using a TOOL/FUNCTION, you MUST WAIT for the tool result before continuing
    with your response. DO NOT MAKE UP RESULTS FROM A TOOL!
[3] YOU MUST NOT ANSWER queries from your OWN KNOWLEDGE; ALWAYS RELY ON 
    the result of a TOOL/FUNCTION to compose your response.
[4] Use ONLY ONE TOOL/FUNCTION at a TIME!
"""
# sys msg to use when schema already provided initially,
# so agent should not use schema tool
SCHEMA_PROVIDED_SYS_MSG = f"""You are a data scientist and expert in Graph Databases, 
with expertise in answering questions by interacting with an ArangoDB database.

The schema below describes the ArangoDB database structure, 
collections (document and edge),
and their attribute keys available in your ArangoDB database.

=== SCHEMA ===
{{schema}}
=== END SCHEMA ===


To help with the user's question or database update/creation request, 
you have access to these tools:

- {aql_retrieval_tool_description}

- {aql_creation_tool_description}


{tool_result_instruction}
"""

# sys msg to use when schema is not initially provided,
# and we want agent to use schema tool to get schema
SCHEMA_TOOLS_SYS_MSG = f"""You are a data scientist and expert in 
Arango Graph Databases, 
with expertise in answering questions by querying ArangoDB database
using the Arango Query Language (AQL).
You have access to the following tools:

- {arango_schema_tool_description}

- {aql_retrieval_tool_description}

- {aql_creation_tool_description}

{tool_result_instruction}
"""

DEFAULT_ARANGO_CHAT_SYSTEM_MESSAGE = f"""
{{mode}}

You do not need to be able to answer a question with just one query. 
You can make a query, WAIT for the result, 
THEN make ANOTHER query, WAIT for result,
THEN make ANOTHER query, and so on, until you have the answer.

{aql_query_instructions}

RETRY-SUGGESTIONS:
If you receive a null or other unexpected result,
(a) make sure you use the available TOOLs correctly,
(b) learn more about the schema using EITHER:
 - `{arango_schema_tool_name}` tool/function-call to find properties of specific
    collections or other parts of the schema, OR
 - `{aql_retrieval_tool_name}` tool/function-call to use AQL queries to 
    find specific parts of the schema.
(c) Collection names are CASE-SENSITIVE -- make sure you adhere to the exact 
    collection name you found in the schema.
(d) see if you have made an assumption in your AQL query, and try another way, 
    or use `{aql_retrieval_tool_name}` to explore the database contents before 
    submitting your final query. 
(e) Try APPROXIMATE or PARTIAL MATCHES to strings in the user's query, 
    e.g. user may ask about "Godfather" instead of "The Godfather",
    or try using CASE-INSENSITIVE MATCHES.
    
Start by asking what the user needs help with.

{tool_result_instruction}

{aql_retrieval_query_example}
"""

ADDRESSING_INSTRUCTION = """
IMPORTANT - Whenever you are NOT writing an AQL query, make sure you address the 
user using {prefix}User. You MUST use the EXACT syntax {prefix} !!!

In other words, you ALWAYS EITHER:
 - write an AQL query using one of the tools, 
 - OR address the user using {prefix}User.
 
YOU CANNOT ADDRESS THE USER WHEN USING A TOOL!!
"""

DONE_INSTRUCTION = f"""
When you are SURE you have the CORRECT answer to a user's query or request, 
use the `{done_tool_name}` with `content` set to the answer or result.
If you DO NOT think you have the answer to the user's query or request,
you SHOULD NOT use the `{done_tool_name}` tool.
Instead, you must CONTINUE to improve your queries (tools) to get the correct answer,
and finally use the `{done_tool_name}` tool to send the correct answer to the user.
"""
</file>

<file path="langroid/agent/special/arangodb/tools.py">
from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage


class AQLRetrievalTool(ToolMessage):
    request: str = "aql_retrieval_tool"
    purpose: str = """
        To send an <aql_query> in response to a user's request/question, 
        OR to find SCHEMA information,
        and WAIT for results of the <aql_query> BEFORE continuing with response.
        You will receive RESULTS from this tool, and ONLY THEN you can continue.
    """
    aql_query: str

    _max_result_tokens: int = 500
    _max_retained_tokens: int = 200

    @classmethod
    def examples(cls) -> List[ToolMessage | Tuple[str, ToolMessage]]:
        """Few-shot examples to include in tool instructions."""
        return [
            (
                "I want to see who Bob's Father is",
                cls(
                    aql_query="""
                    FOR v, e, p IN 1..1 OUTBOUND 'users/Bob' GRAPH 'family_tree'
                    FILTER p.edges[0].type == 'father'
                    RETURN v
                    """
                ),
            ),
            (
                "I want to know the properties of the Actor node",
                cls(
                    aql_query="""
                    FOR doc IN Actor
                    LIMIT 1
                    RETURN ATTRIBUTES(doc)                    
                    """
                ),
            ),
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        When using this TOOL/Function-call, you must WAIT to receive the RESULTS 
        of the AQL query, before continuing your response!
        DO NOT ASSUME YOU KNOW THE RESULTs BEFORE RECEIVING THEM.        
        """


aql_retrieval_tool_name = AQLRetrievalTool.default_value("request")


class AQLCreationTool(ToolMessage):
    request: str = "aql_creation_tool"
    purpose: str = """
        To send the <aql_query> to create documents/edges in the graph database.
        IMPORTANT: YOU MUST WAIT FOR THE RESULT OF THE TOOL BEFORE CONTINUING.
        You will receive RESULTS from this tool, and ONLY THEN you can continue.
    """
    aql_query: str

    @classmethod
    def examples(cls) -> List[ToolMessage | Tuple[str, ToolMessage]]:
        """Few-shot examples to include in tool instructions."""
        return [
            (
                "Create a new document in the collection 'users'",
                cls(
                    aql_query="""
                    INSERT {
                      "name": "Alice",
                      "age": 30
                    } INTO users
                    """
                ),
            ),
        ]


aql_creation_tool_name = AQLCreationTool.default_value("request")


class ArangoSchemaTool(ToolMessage):
    request: str = "arango_schema_tool"
    purpose: str = """
        To get the schema of the Arango graph database,
        or some part of it. Follow these instructions:
        1. Set <properties> to True to get the properties of the collections,
        and False if you only want to see the graph structure and get only the
        from/to relations of the edges.
        2. Set <collections> to a list of collection names if you want to see,
        or leave it as None to see all ALL collections.
        IMPORTANT: YOU MUST WAIT FOR THE RESULT OF THE TOOL BEFORE CONTINUING.
        You will receive RESULTS from this tool, and ONLY THEN you can continue.
    """

    properties: bool = True
    collections: List[str] | None = None

    _max_result_tokens: int = 500


arango_schema_tool_name = ArangoSchemaTool.default_value("request")
</file>

<file path="langroid/agent/special/arangodb/utils.py">
from typing import Any, Dict, List


def count_fields(schema: Dict[str, List[Dict[str, Any]]]) -> int:
    total = 0
    for coll in schema["Collection Schema"]:
        # Count all keys in each collection's dict
        total += len(coll)
        # Also count properties if they exist
        props = coll.get(f"{coll['collection_type']}_properties", [])
        total += len(props)
    return total


def trim_schema(
    schema: Dict[str, List[Dict[str, Any]]]
) -> Dict[str, List[Dict[str, Any]]]:
    """Keep only edge connection info, remove properties and examples"""
    trimmed: Dict[str, List[Dict[str, Any]]] = {
        "Graph Schema": schema["Graph Schema"],
        "Collection Schema": [],
    }
    for coll in schema["Collection Schema"]:
        col_info: Dict[str, Any] = {
            "collection_name": coll["collection_name"],
            "collection_type": coll["collection_type"],
        }
        if coll["collection_type"] == "edge":
            # preserve from/to info if present
            if f"example_{coll['collection_type']}" in coll:
                example = coll[f"example_{coll['collection_type']}"]
                if example and "_from" in example:
                    col_info["from_collection"] = example["_from"].split("/")[0]
                    col_info["to_collection"] = example["_to"].split("/")[0]
        trimmed["Collection Schema"].append(col_info)
    return trimmed
</file>

<file path="langroid/agent/special/lance_rag/__init__.py">
from . import query_planner_agent
from . import critic_agent
from . import lance_rag_task

__all__ = [
    "query_planner_agent",
    "critic_agent",
    "lance_rag_task",
]
</file>

<file path="langroid/agent/special/lance_rag/critic_agent.py">
"""
QueryPlanCritic is a ChatAgent that is created with a specific document schema.

Its role is to provide feedback on a Query Plan, which consists of:
- filter condition if needed (or empty string if no filter is needed)
- query - a possibly rephrased query that can be used to match the `content` field
- dataframe_calc - a Pandas-dataframe calculation/aggregation string, possibly empty
- original_query - the original query for reference
- result - the answer received from an assistant that used this QUERY PLAN.

This agent has access to two tools:
- QueryPlanTool: The handler method for this tool re-writes the query plan
  in plain text (non-JSON) so the LLM can provide its feedback using the
  QueryPlanFeedbackTool.
- QueryPlanFeedbackTool: LLM uses this tool to provide feedback on the Query Plan
"""

import logging

from langroid.agent.chat_agent import ChatAgent
from langroid.agent.chat_document import ChatDocument
from langroid.agent.special.lance_rag.query_planner_agent import (
    LanceQueryPlanAgentConfig,
)
from langroid.agent.special.lance_tools import (
    QueryPlanAnswerTool,
    QueryPlanFeedbackTool,
)
from langroid.agent.tools.orchestration import AgentDoneTool
from langroid.utils.constants import NO_ANSWER

logger = logging.getLogger(__name__)


class QueryPlanCriticConfig(LanceQueryPlanAgentConfig):
    name: str = "QueryPlanCritic"
    system_message: str = f"""
    You are an expert at carefully planning a query that needs to be answered
    based on a large collection of documents. These docs have a special `content` field
    and additional FILTERABLE fields in the SCHEMA below, along with the 
    SAMPLE VALUES for each field, and the DTYPE in PANDAS TERMINOLOGY.
    
    {{doc_schema}}
    
    The ORIGINAL QUERY is handled by a QUERY PLANNER who sends the PLAN to an ASSISTANT,
    who returns an ANSWER.
    
    You will receive a QUERY PLAN consisting of:
    - ORIGINAL QUERY from the user, which a QUERY PLANNER processes,
      to create a QUERY PLAN, to be handled by an ASSISTANT.
    - PANDAS-LIKE FILTER, WHICH CAN BE EMPTY (and it's fine if results sound reasonable)
      FILTER SHOULD ONLY BE USED IF EXPLICITLY REQUIRED BY THE QUERY.
      This filter selects the documents over which the REPHRASED QUERY will be applied,
      thus naturally, the Re-phrased Query should NOT mention any FILTER fields,
      since it applies to the documents AFTER FILTERING.
    - REPHRASED QUERY (CANNOT BE EMPTY) that will be used to match against the 
      CONTENT (not filterable) of the documents.
      In general the REPHRASED QUERY should be relied upon to match the CONTENT 
      of the docs. Thus the REPHRASED QUERY itself acts like a 
      SEMANTIC/LEXICAL/FUZZY FILTER since the Assistant is able to use it to match 
      the CONTENT of the docs in various ways (semantic, lexical, fuzzy, etc.). 
        Keep in mind that the ASSISTANT does NOT know anything about the FILTER fields,
        so the REPHRASED QUERY should NOT mention ANY FILTER fields.
        The assistant will answer based on documents whose CONTENTS match the QUERY, 
        possibly REPHRASED. 
        !!!!****THE REPHRASED QUERY SHOULD NEVER BE EMPTY****!!!
        
        
    - DATAFRAME CALCULATION, which must be a SINGLE LINE calculation (or empty),
        [NOTE ==> This calculation is applied AFTER the FILTER and REPHRASED QUERY.],
    - ANSWER received from an assistant that used this QUERY PLAN.
      IT IS TOTALLY FINE FOR THE ANSWER TO NOT MENTION ANY FILTERING CONDITIONS,
      or if the ANSWER STATEMENT is MISSING SOME CRITERIA in the ORIGINAL QUERY.

        Here is an example of a VALID Plan + Answer:
        
        ORIGINAL QUERY: "Which crime novels were written by Russian authors after 1900?"
        FILTER: "author_nationality == 'Russian' and year_written > 1900"
        REPHRASED QUERY: "crime novel" [NOTICE NO FILTER FIELDS MENTIONED!!!]
        DATAFRAME CALC: ""
        ANSWER: "The Master and Margarita by Mikhail Bulgakov" 
            [NOTICE the answer does NOT need to say "crime novel" or "russian author"]
            
            
        Other examples of VALID ANSWER for a given ORIGINAL QUERY:
        
        ORIGINAL QUERY: "Which mountain is taller than 8000 meters?"
        ANSWER: "Mount Everest" [NOTICE no mention of "taller than 8000 meters"]
        
        ORIGINAL QUERY: "Which country has hosted the most olympics?"
        ANSWER: "United States" [NOTICE no mention of "most olympics"]

    In addition to the above SCHEMA fields there is a `content` field which:
    - CANNOT appear in a FILTER, 
    - CAN appear in the DATAFRAME CALCULATION.
    THERE ARE NO OTHER FIELDS IN THE DOCUMENTS or in the RESULTING DATAFRAME.
        
    Your job is to act as a CRITIC and provide feedback, 
    ONLY using the `query_plan_feedback` tool, and DO NOT SAY ANYTHING ELSE.
    
    Here is how you must examine the QUERY PLAN + ANSWER:
    - ALL filtering conditions in the original query must be EXPLICITLY 
      mentioned in the FILTER, and the QUERY field should not be used for filtering.
    - If the ANSWER contains an ERROR message, then this means that the query
      plan execution FAILED, and your feedback should say INVALID along 
      with the ERROR message, `suggested_fix` that aims to help the assistant 
      fix the problem (or simply equals "address the the error shown in feedback")
    - Ask yourself, is the ANSWER in the expected form, e.g. 
        if the question is asking for the name of an ENTITY with max SIZE,
        then the answer should be the ENTITY name, NOT the SIZE!! 
    - If the ANSWER is in the expected form, then the QUERY PLAN is likely VALID,
      and your feedback should say VALID, with empty `suggested_fix`.
      ===> HOWEVER!!! Watch out for a spurious correct-looking answer, for EXAMPLE:
      the query was to find the ENTITY with a maximum SIZE, 
      but the dataframe calculation is find the SIZE, NOT the ENTITY!!      
    - If the ANSWER is {NO_ANSWER} or of the wrong form, 
      then try to DIAGNOSE the problem IN THE FOLLOWING ORDER:
      - DATAFRAME CALCULATION -- is it doing the right thing?
        Is it finding the Index of a row instead of the value in a column?
        Or another example: maybe it is finding the maximum population
           rather than the CITY with the maximum population?
        If you notice a problem with the DATAFRAME CALCULATION, then
        ONLY SUBMIT FEEDBACK ON THE DATAFRAME CALCULATION, and DO NOT
        SUGGEST ANYTHING ELSE.
      - If the DATAFRAME CALCULATION looks correct, then check if 
        the REPHRASED QUERY makes sense given the ORIGINAL QUERY and FILTER.
        If this is the problem, then ONLY SUBMIT FEEDBACK ON THE REPHRASED QUERY,
        and DO NOT SUGGEST ANYTHING ELSE.
      - If the REPHRASED QUERY looks correct, then check if the FILTER makes sense.
        REMEMBER: A filter should ONLY be used if EXPLICITLY REQUIRED BY THE QUERY.
     
     
     IMPORTANT!! The DATAFRAME CALCULATION is done AFTER applying the 
         FILTER and REPHRASED QUERY! Keep this in mind when evaluating 
         the correctness of the DATAFRAME CALCULATION.
    
    ALWAYS use `query_plan_feedback` tool/fn to present your feedback
    in the `feedback` field, and if any fix is suggested,
    present it in the `suggested_fix` field.
    DO NOT SAY ANYTHING ELSE OUTSIDE THE TOOL/FN.
    IF NO REVISION NEEDED, simply leave the `suggested_fix` field EMPTY,
    and SAY NOTHING ELSE
    and DO NOT EXPLAIN YOURSELF.        
    """


def plain_text_query_plan(msg: QueryPlanAnswerTool) -> str:
    plan = f"""
    OriginalQuery: {msg.plan.original_query}
    Filter: {msg.plan.filter}
    Rephrased Query: {msg.plan.query}
    DataframeCalc: {msg.plan.dataframe_calc}
    Answer: {msg.answer}
    """
    return plan


class QueryPlanCritic(ChatAgent):
    """
    Critic for LanceQueryPlanAgent, provides feedback on
    query plan + answer.
    """

    def __init__(self, cfg: LanceQueryPlanAgentConfig):
        super().__init__(cfg)
        self.config = cfg
        self.enable_message(QueryPlanAnswerTool, use=False, handle=True)
        self.enable_message(QueryPlanFeedbackTool, use=True, handle=True)
        self.enable_message(AgentDoneTool, use=False, handle=True)

    def init_state(self) -> None:
        super().init_state()
        self.expecting_feedback_tool = False

    def query_plan_answer(self, msg: QueryPlanAnswerTool) -> str:
        """Present query plan + answer in plain text (not JSON)
        so LLM can give feedback"""
        self.expecting_feedback_tool = True
        return plain_text_query_plan(msg)

    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> AgentDoneTool:
        """Format Valid so return to Query Planner"""
        self.expecting_feedback_tool = False
        # indicate this task is Done, and return the tool as result
        return AgentDoneTool(tools=[msg])

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        """Remind the LLM to use QueryPlanFeedbackTool since it forgot"""
        if self.expecting_feedback_tool:
            return """
            You forgot to use the `query_plan_feedback` tool/function.
            Re-try your response using the `query_plan_feedback` tool/function,
            remember to provide feedback in the `feedback` field,
            and if any fix is suggested, provide it in the `suggested_fix` field.
            """
        return None
</file>

<file path="langroid/agent/special/lance_rag/lance_rag_task.py">
"""
The LanceRAGTaskCreator.new() method creates a 3-Agent system that uses this agent.
It takes a LanceDocChatAgent instance as argument, and adds two more agents:
- LanceQueryPlanAgent, which is given the LanceDB schema in LanceDocChatAgent,
and based on this schema, for a given user query, creates a Query Plan
using the QueryPlanTool, which contains a filter, a rephrased query,
and a dataframe_calc.
- QueryPlanCritic, which is given the LanceDB schema in LanceDocChatAgent,
 and gives feedback on the Query Plan and Result using the QueryPlanFeedbackTool.

The LanceRAGTaskCreator.new() method sets up the given LanceDocChatAgent and
QueryPlanCritic as sub-tasks of the LanceQueryPlanAgent's task.

Langroid's built-in task orchestration ensures that:
- the LanceQueryPlanAgent reformulates the plan based
    on the QueryPlanCritics's feedback,
- LLM deviations are corrected via tools and overrides of ChatAgent methods.
"""

import logging

from langroid.agent.special.lance_doc_chat_agent import LanceDocChatAgent
from langroid.agent.special.lance_rag.critic_agent import (
    QueryPlanCritic,
    QueryPlanCriticConfig,
)
from langroid.agent.special.lance_rag.query_planner_agent import (
    LanceQueryPlanAgent,
    LanceQueryPlanAgentConfig,
)
from langroid.agent.task import Task
from langroid.mytypes import Entity

logger = logging.getLogger(__name__)


class LanceRAGTaskCreator:
    @staticmethod
    def new(
        agent: LanceDocChatAgent,
        interactive: bool = True,
    ) -> Task:
        """
        Add a LanceFilterAgent to the LanceDocChatAgent,
        set up the corresponding Tasks, connect them,
        and return the top-level query_plan_task.
        """
        doc_agent_name = "LanceRAG"
        critic_name = "QueryPlanCritic"
        query_plan_agent_config = LanceQueryPlanAgentConfig(
            critic_name=critic_name,
            doc_agent_name=doc_agent_name,
            doc_schema=agent._get_clean_vecdb_schema(),
            llm=agent.config.llm,
        )
        query_plan_agent_config.set_system_message()

        critic_config = QueryPlanCriticConfig(
            doc_schema=agent._get_clean_vecdb_schema(),
            llm=agent.config.llm,
        )
        critic_config.set_system_message()

        query_planner = LanceQueryPlanAgent(query_plan_agent_config)
        query_plan_task = Task(
            query_planner,
            interactive=interactive,
        )
        critic_agent = QueryPlanCritic(critic_config)
        critic_task = Task(
            critic_agent,
            interactive=False,
        )
        rag_task = Task(
            agent,
            name="LanceRAG",
            interactive=False,
            done_if_response=[Entity.LLM],  # done when non-null response from LLM
            done_if_no_response=[Entity.LLM],  # done when null response from LLM
        )
        query_plan_task.add_sub_task([critic_task, rag_task])
        return query_plan_task
</file>

<file path="langroid/agent/special/lance_rag/query_planner_agent.py">
"""
LanceQueryPlanAgent is a ChatAgent created with a specific document schema.
Given a QUERY, the LLM constructs a Query Plan consisting of:
- filter condition if needed (or empty string if no filter is needed)
- query - a possibly rephrased query that can be used to match the `content` field
- dataframe_calc - a Pandas-dataframe calculation/aggregation string, possibly empty
- original_query - the original query for reference

This agent has access to two tools:
- QueryPlanTool, which is used to generate the Query Plan, and the handler of
    this tool simply passes it on to the RAG agent named in config.doc_agent_name.
- QueryPlanFeedbackTool, which is used to handle feedback on the Query Plan and
  Result from the RAG agent. The QueryPlanFeedbackTool is used by
  the QueryPlanCritic, who inserts feedback into the `feedback` field
"""

import logging
from typing import Optional

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.chat_document import ChatDocument
from langroid.agent.special.lance_tools import (
    AnswerTool,
    QueryPlan,
    QueryPlanAnswerTool,
    QueryPlanFeedbackTool,
    QueryPlanTool,
)
from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool
from langroid.utils.constants import NO_ANSWER

logger = logging.getLogger(__name__)


class LanceQueryPlanAgentConfig(ChatAgentConfig):
    name: str = "LancePlanner"
    critic_name: str = "QueryPlanCritic"
    doc_agent_name: str = "LanceRAG"
    doc_schema: str = ""
    use_tools: bool = False
    max_retries: int = 5  # max number of retries for query plan
    use_functions_api: bool = True

    system_message: str = """
    You will receive a QUERY, to be answered based on an EXTREMELY LARGE collection
    of documents you DO NOT have access to, but your ASSISTANT does.
    You only know that these documents have a special `content` field
    and additional FILTERABLE fields in the SCHEMA below, along with the 
    SAMPLE VALUES for each field, and the DTYPE in PANDAS TERMINOLOGY.
    
    {doc_schema}
    
    Based on the QUERY and the above SCHEMA, your task is to determine a QUERY PLAN,
    consisting of:
    -  a PANDAS-TYPE FILTER (can be empty string) that would help the ASSISTANT to 
        answer the query.
        Remember the FILTER can refer to ANY fields in the above SCHEMA
        EXCEPT the `content` field of the documents. 
        ONLY USE A FILTER IF EXPLICITLY MENTIONED IN THE QUERY.
        TO get good results, for STRING MATCHES, consider using LIKE instead of =, e.g.
        "CEO LIKE '%Jobs%'" instead of "CEO = 'Steve Jobs'"
        YOUR FILTER MUST BE A PANDAS-TYPE FILTER, respecting the shown DTYPES.        
    - a possibly REPHRASED QUERY (CANNOT BE EMPTY) to be answerable given the FILTER.
        Keep in mind that the ASSISTANT does NOT know anything about the FILTER fields,
        so the REPHRASED QUERY should NOT mention ANY FILTER fields.
        The assistant will answer based on documents whose CONTENTS match the QUERY, 
        possibly REPHRASED. 
        !!!!****THE REPHRASED QUERY SHOULD NEVER BE EMPTY****!!!
    - an OPTIONAL SINGLE-LINE Pandas-dataframe calculation/aggregation string 
        that can be used to calculate the answer to the original query, 
        e.g. "df["rating"].mean()",
        or "df.groupby("director").mean()["rating"]", 
        or EMPTY string if no calc is needed. 
        The dataframe calc CAN refer to the `content` field.
        If a DataFrame calculation is NOT needed, leave this field EMPTY.
        
        IMPORTANT: The DataFrame `df` in this calculation is the result of 
        applying the FILTER AND REPHRASED QUERY to the documents.
        
        WATCH OUT!! When deciding the dataframe calc, if any, CAREFULLY
        note what the query is asking, and ensure that the result of your
        dataframe calc expression would answer the query.                
    
    
    EXAMPLE:
    ------- 
    Suppose there is a document-set about crime reports, where:
     CONTENT = crime report,
     Filterable SCHEMA consists of City, Year, num_deaths.
    
    Then given this ORIGINAL QUERY: 
    
        Total deaths in shoplifting crimes in Los Angeles in 2023?
    
    A POSSIBLE QUERY PLAN could be:
    
    FILTER: "City LIKE '%Los Angeles%' AND Year = 2023"
    REPHRASED QUERY: "shoplifting crime" --> this will be used to MATCH content of docs
         [NOTE: we dropped the FILTER fields City and Year since the 
         ASSISTANT does not know about them and only uses the query to 
         match the CONTENT of the docs.]
    DATAFRAME CALCULATION: "df["num_deaths"].sum()"
        NOTE!!! The DataFrame `df` in this calculation is the result of
        applying the FILTER AND REPHRASED QUERY to the documents, 
        hence this computation will give the total deaths in shoplifting crimes.
    ------------- END OF EXAMPLE ----------------
    
    The FILTER must be a PANDAS-like condition, e.g. 
    "year > 2000 AND genre = 'ScienceFiction'".
    To ensure you get useful results, you should make your FILTER 
    NOT TOO STRICT, e.g. look for approximate match using LIKE, etc.
    E.g. "CEO LIKE '%Jobs%'" instead of "CEO = 'Steve Jobs'"
    Use DOT NOTATION to refer to nested fields, e.g. `metadata.year`, etc. 
        
    You must FIRST present the QUERY PLAN using the `query_plan` tool/function.
    This will be handled by your document assistant, who will produce an ANSWER.
            
    You may receive FEEDBACK on your QUERY PLAN and received ANSWER,
    from the 'QueryPlanCritic' who may offer suggestions for
    a better FILTER, REPHRASED QUERY, or DATAFRAME CALCULATION.
                  
    At the BEGINNING if there is no query, ASK the user what they want to know.
    """

    def set_system_message(self) -> None:
        self.system_message = self.system_message.format(
            doc_schema=self.doc_schema,
        )


class LanceQueryPlanAgent(ChatAgent):
    def __init__(self, config: LanceQueryPlanAgentConfig):
        super().__init__(config)
        self.config: LanceQueryPlanAgentConfig = config
        # This agent should generate the QueryPlanTool
        # as well as handle it for validation
        self.enable_message(QueryPlanTool, use=True, handle=True)
        self.enable_message(QueryPlanFeedbackTool, use=False, handle=True)
        self.enable_message(AnswerTool, use=False, handle=True)
        # neither use nor handle! Added to "known" tools so that the Planner agent
        # can avoid processing it
        self.enable_message(QueryPlanAnswerTool, use=False, handle=False)
        # LLM will not use this, so set use=False (Agent generates it)
        self.enable_message(AgentDoneTool, use=False, handle=True)

    def init_state(self) -> None:
        super().init_state()
        self.curr_query_plan: QueryPlan | None = None
        self.expecting_query_plan: bool = False
        # how many times re-trying query plan in response to feedback:
        self.n_retries: int = 0
        self.n_query_plan_reminders: int = 0
        self.result: str = ""  # answer received from LanceRAG

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        self.expecting_query_plan = True
        return super().llm_response(message)

    def query_plan(self, msg: QueryPlanTool) -> ForwardTool | str:
        """Valid, tool msg, forward chat_doc to RAG Agent.
        Note this chat_doc will already have the
        QueryPlanTool in its tool_messages list.
        We just update the recipient to the doc_agent_name.
        """
        # save, to be used to assemble QueryPlanResultTool
        if len(msg.plan.dataframe_calc.split("\n")) > 1:
            return "DATAFRAME CALCULATION must be a SINGLE LINE; Retry the `query_plan`"
        self.curr_query_plan = msg.plan
        self.expecting_query_plan = False

        # To forward the QueryPlanTool to doc_agent, we could either:

        # (a) insert `recipient` in the QueryPlanTool:
        # QPWithRecipient = QueryPlanTool.require_recipient()
        # qp = QPWithRecipient(**msg.model_dump(), recipient=self.config.doc_agent_name)
        # return qp
        #
        # OR
        #
        # (b) create an agent response with recipient and tool_messages.
        # response = self.create_agent_response(
        #     recipient=self.config.doc_agent_name, tool_messages=[msg]
        # )
        # return response

        # OR
        # (c) use the ForwardTool:
        return ForwardTool(agent=self.config.doc_agent_name)

    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> str | AgentDoneTool:
        """Process Critic feedback on QueryPlan + Answer from RAG Agent"""
        # We should have saved answer in self.result by this time,
        # since this Agent seeks feedback only after receiving RAG answer.
        if (
            msg.suggested_fix == ""
            and NO_ANSWER not in self.result
            and self.result != ""
        ):
            # This means the result is good AND Query Plan is fine,
            # as judged by Critic
            # (Note sometimes critic may have empty suggested_fix even when
            # the result is NO_ANSWER)
            self.n_retries = 0  # good answer, so reset this
            return AgentDoneTool(content=self.result)
        self.n_retries += 1
        if self.n_retries >= self.config.max_retries:
            # bail out to avoid infinite loop
            self.n_retries = 0
            return AgentDoneTool(content=NO_ANSWER)

        # there is a suggested_fix, OR the result is empty or NO_ANSWER
        if self.result == "" or NO_ANSWER in self.result:
            # if result is empty or NO_ANSWER, we should retry the query plan
            feedback = """
            There was no answer, which might mean there is a problem in your query.
            """
            suggested = "Retry the `query_plan` to try to get a non-null answer"
        else:
            feedback = msg.feedback
            suggested = msg.suggested_fix

        self.expecting_query_plan = True

        return f"""
        here is FEEDBACK about your QUERY PLAN, and a SUGGESTED FIX.
        Modify the QUERY PLAN if needed:
        ANSWER: {self.result}
        FEEDBACK: {feedback}
        SUGGESTED FIX: {suggested}
        """

    def answer_tool(self, msg: AnswerTool) -> QueryPlanAnswerTool:
        """Handle AnswerTool received from LanceRagAgent:
        Construct a QueryPlanAnswerTool with the answer"""
        self.result = msg.answer  # save answer to interpret feedback later
        assert self.curr_query_plan is not None
        query_plan_answer_tool = QueryPlanAnswerTool(
            plan=self.curr_query_plan,
            answer=msg.answer,
        )
        self.curr_query_plan = None  # reset
        return query_plan_answer_tool

    def handle_message_fallback(
        self, msg: str | ChatDocument
    ) -> str | ChatDocument | None:
        """
        Remind to use QueryPlanTool if we are expecting it.
        """
        if self.expecting_query_plan and self.n_query_plan_reminders < 5:
            self.n_query_plan_reminders += 1
            return """
            You FORGOT to use the `query_plan` tool/function, 
            OR you had a WRONG JSON SYNTAX when trying to use it.
            Re-try your response using the `query_plan` tool/function CORRECTLY.
            """
        self.n_query_plan_reminders = 0  # reset
        return None
</file>

<file path="langroid/agent/special/neo4j/__init__.py">

</file>

<file path="langroid/agent/special/neo4j/system_messages.py">
from langroid.agent.special.neo4j.tools import (
    cypher_creation_tool_name,
    cypher_retrieval_tool_name,
    graph_schema_tool_name,
)
from langroid.agent.tools.orchestration import DoneTool

done_tool_name = DoneTool.default_value("request")

graph_schema_tool_description = f"""
`{graph_schema_tool_name}` tool/function-call to get all the node labels, relationship 
 types, and property keys available in your Neo4j database. You MUST use
 this tool BEFORE attempting to use the `{cypher_retrieval_tool_name}` tool,
 to ensure that you are using the correct node labels, relationship types, and
 property keys in your `{cypher_retrieval_tool_name}` tool/function-call.
"""

cypher_retrieval_tool_description = f"""
`{cypher_retrieval_tool_name}` tool/function-call to retrieve information from the 
     graph database to answer questions.
"""

cypher_creation_tool_description = f"""
`{cypher_creation_tool_name}` tool/function-call to execute cypher query that creates
   entities/relationships in the graph database.
"""

cypher_query_instructions = """
You must be smart about using the right node labels, relationship types, and property
keys based on the english description. If you are thinking of using a node label,
relationship type, or property key that does not exist, you are probably on the wrong 
track, so you should try your best to answer based on an existing table or column.
DO NOT assume any nodes or relationships other than those above.
"""


# sys msg to use when schema already provided initially,
# so agent does not need to use schema tool, at least initially,
# but may do so later if the db evolves, or if needs to bring in the schema
# to more recent context.
SCHEMA_PROVIDED_SYS_MSG = f"""You are a data scientist and expert in Knowledge Graphs, 
with expertise in answering questions by interacting with a Neo4j graph database.

The schema below describes the Neo4j database structure, node labels, 
relationship types, and property keys available in your Neo4j database.

=== SCHEMA ===
{{schema}}
=== END SCHEMA ===

To help with the user's question or database update/creation request, 
you have access to these tools:

- {cypher_retrieval_tool_description}

- {cypher_creation_tool_description}

Since the schema has been provided, you may not need to use the tool below,
but you may use it if you need to remind yourself about the schema:

- {graph_schema_tool_description}
 
"""

# sys msg to use when schema is not initially provided,
# and we want agent to use schema tool to get schema
SCHEMA_TOOLS_SYS_MSG = f"""You are a data scientist and expert in Knowledge Graphs, 
with expertise in answering questions by querying Neo4j database.
You have access to the following tools:

- {graph_schema_tool_description}

- {cypher_retrieval_tool_description}

- {cypher_creation_tool_description}

"""

DEFAULT_NEO4J_CHAT_SYSTEM_MESSAGE = f"""
{{mode}}

You do not need to be able to answer a question with just one query. 
You could make a sequence of Cypher queries to find the answer to the question.

{cypher_query_instructions}



RETRY-SUGGESTIONS:
If you receive a null or other unexpected result,
(a) make sure you use the available TOOLs correctly,
(b) USE `{graph_schema_tool_name}` tool/function-call to get all the node labels, 
    relationship types, and property keys available in your Neo4j database. 
(c) LABELS are CASE-SENSITIVE -- make sure you adhere to the exact label name
   you found in the schema.
(d) see if you have made an assumption in your Neo4j query, and try another way, 
   or use `{cypher_retrieval_tool_name}` to explore the database contents before 
   submitting your final query. 
(e) USE `{cypher_creation_tool_name}` tool/function-call to execute cypher query that 
    creates entities/relationships in the graph database.
(f) Try APPROXIMATE or PARTIAL MATCHES to strings in the user's query, 
    e.g. user may ask about "Godfather" instead of "The Godfather",
    or try using CASE-INSENSITIVE MATCHES.

Start by asking what the user needs help with.
"""

ADDRESSING_INSTRUCTION = """
IMPORTANT - Whenever you are NOT writing a CYPHER query, make sure you address the 
user using {prefix}User. You MUST use the EXACT syntax {prefix} !!!

In other words, you ALWAYS EITHER:
 - write a CYPHER query using one of the tools, 
 - OR address the user using {prefix}User.
"""

DONE_INSTRUCTION = f"""
When you finally have the answer to a user's query or request, 
use the `{done_tool_name}` with `content` set to the answer or result.
"""
</file>

<file path="langroid/agent/special/neo4j/tools.py">
from langroid.agent import ToolMessage


class CypherRetrievalTool(ToolMessage):
    request: str = "cypher_retrieval_tool"
    purpose: str = """To send the <cypher_query> to retrieve 
        data from the graph database based on provided text description and schema.
        """
    cypher_query: str


cypher_retrieval_tool_name = CypherRetrievalTool.default_value("request")


class CypherCreationTool(ToolMessage):
    request: str = "cypher_creation_tool"
    purpose: str = """
        To send the <cypher_query> to create 
        entities/relationships in the graph database.
        """
    cypher_query: str


cypher_creation_tool_name = CypherCreationTool.default_value("request")


class GraphSchemaTool(ToolMessage):
    request: str = "graph_schema_tool"
    purpose: str = """To get the schema of the graph database."""


graph_schema_tool_name = GraphSchemaTool.default_value("request")
</file>

<file path="langroid/agent/special/sql/utils/__init__.py">
from . import tools
from . import description_extractors
from . import populate_metadata
from . import system_message
from .tools import (
    RunQueryTool,
    GetTableNamesTool,
    GetTableSchemaTool,
    GetColumnDescriptionsTool,
)

__all__ = [
    "RunQueryTool",
    "GetTableNamesTool",
    "GetTableSchemaTool",
    "GetColumnDescriptionsTool",
    "description_extractors",
    "populate_metadata",
    "system_message",
    "tools",
]
</file>

<file path="langroid/agent/special/sql/utils/description_extractors.py">
from typing import Any, Dict, List, Optional

from langroid.exceptions import LangroidImportError

try:
    from sqlalchemy import inspect, text
    from sqlalchemy.engine import Engine
except ImportError as e:
    raise LangroidImportError(extra="sql", error=str(e))


def extract_postgresql_descriptions(
    engine: Engine,
    multi_schema: bool = False,
) -> Dict[str, Dict[str, Any]]:
    """
    Extracts descriptions for tables and columns from a PostgreSQL database.

    This method retrieves the descriptions of tables and their columns
    from a PostgreSQL database using the provided SQLAlchemy engine.

    Args:
        engine (Engine): SQLAlchemy engine connected to a PostgreSQL database.
        multi_schema (bool): Generate descriptions for all schemas in the database.

    Returns:
        Dict[str, Dict[str, Any]]: A dictionary mapping table names to a
        dictionary containing the table description and a dictionary of
        column descriptions.
    """
    inspector = inspect(engine)
    result: Dict[str, Dict[str, Any]] = {}

    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
        table_names: List[str] = inspector.get_table_names(schema=schema)
        with engine.connect() as conn:
            for table in table_names:
                if schema is None:
                    table_name = table
                else:
                    table_name = f"{schema}.{table}"

                table_comment = (
                    conn.execute(
                        text(f"SELECT obj_description('{table_name}'::regclass)")
                    ).scalar()
                    or ""
                )

                columns = {}
                col_data = inspector.get_columns(table, schema=schema)
                for idx, col in enumerate(col_data, start=1):
                    col_comment = (
                        conn.execute(
                            text(
                                f"SELECT col_description('{table_name}'::regclass, "
                                f"{idx})"
                            )
                        ).scalar()
                        or ""
                    )
                    columns[col["name"]] = col_comment

                result[table_name] = {"description": table_comment, "columns": columns}

    if multi_schema:
        for schema in inspector.get_schema_names():
            gen_schema_descriptions(schema)
    else:
        gen_schema_descriptions()

    return result


def extract_mysql_descriptions(
    engine: Engine,
    multi_schema: bool = False,
) -> Dict[str, Dict[str, Any]]:
    """Extracts descriptions for tables and columns from a MySQL database.

    This method retrieves the descriptions of tables and their columns
    from a MySQL database using the provided SQLAlchemy engine.

    Args:
        engine (Engine): SQLAlchemy engine connected to a MySQL database.
        multi_schema (bool): Generate descriptions for all schemas in the database.

    Returns:
        Dict[str, Dict[str, Any]]: A dictionary mapping table names to a
        dictionary containing the table description and a dictionary of
        column descriptions.
    """
    inspector = inspect(engine)
    result: Dict[str, Dict[str, Any]] = {}

    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
        table_names: List[str] = inspector.get_table_names(schema=schema)

        with engine.connect() as conn:
            for table in table_names:
                if schema is None:
                    table_name = table
                else:
                    table_name = f"{schema}.{table}"

                query = text(
                    "SELECT table_comment FROM information_schema.tables WHERE"
                    " table_schema = :schema AND table_name = :table"
                )
                table_result = conn.execute(
                    query, {"schema": engine.url.database, "table": table_name}
                )
                table_comment = table_result.scalar() or ""

                columns = {}
                for col in inspector.get_columns(table, schema=schema):
                    columns[col["name"]] = col.get("comment", "")

                result[table_name] = {"description": table_comment, "columns": columns}

    if multi_schema:
        for schema in inspector.get_schema_names():
            gen_schema_descriptions(schema)
    else:
        gen_schema_descriptions()

    return result


def extract_default_descriptions(
    engine: Engine, multi_schema: bool = False
) -> Dict[str, Dict[str, Any]]:
    """Extracts default descriptions for tables and columns from a database.

    This method retrieves the table and column names from the given database
    and associates empty descriptions with them.

    Args:
        engine (Engine): SQLAlchemy engine connected to a database.
        multi_schema (bool): Generate descriptions for all schemas in the database.

    Returns:
        Dict[str, Dict[str, Any]]: A dictionary mapping table names to a
        dictionary containing an empty table description and a dictionary of
        empty column descriptions.
    """
    inspector = inspect(engine)
    result: Dict[str, Dict[str, Any]] = {}

    def gen_schema_descriptions(schema: Optional[str] = None) -> None:
        table_names: List[str] = inspector.get_table_names(schema=schema)

        for table in table_names:
            columns = {}
            for col in inspector.get_columns(table):
                columns[col["name"]] = ""

            result[table] = {"description": "", "columns": columns}

    if multi_schema:
        for schema in inspector.get_schema_names():
            gen_schema_descriptions(schema)
    else:
        gen_schema_descriptions()

    return result


def extract_schema_descriptions(
    engine: Engine, multi_schema: bool = False
) -> Dict[str, Dict[str, Any]]:
    """
    Extracts the schema descriptions from the database connected to by the engine.

    Args:
        engine (Engine): SQLAlchemy engine instance.
        multi_schema (bool): Generate descriptions for all schemas in the database.

    Returns:
        Dict[str, Dict[str, Any]]: A dictionary representation of table and column
        descriptions.
    """

    extractors = {
        "postgresql": extract_postgresql_descriptions,
        "mysql": extract_mysql_descriptions,
    }
    return extractors.get(engine.dialect.name, extract_default_descriptions)(
        engine, multi_schema=multi_schema
    )
</file>

<file path="langroid/agent/special/sql/utils/populate_metadata.py">
from typing import Dict, List, Union

from langroid.exceptions import LangroidImportError

try:
    from sqlalchemy import MetaData
except ImportError as e:
    raise LangroidImportError(extra="sql", error=str(e))


def populate_metadata_with_schema_tools(
    metadata: MetaData | List[MetaData],
    info: Dict[str, Dict[str, Union[str, Dict[str, str]]]],
) -> Dict[str, Dict[str, Union[str, Dict[str, str]]]]:
    """
    Extracts information from an SQLAlchemy database's metadata and combines it
    with another dictionary with context descriptions.

    Args:
        metadata (MetaData): SQLAlchemy metadata object of the database.
        info (Dict[str, Dict[str, Any]]): A dictionary with table and column
                                             descriptions.

    Returns:
        Dict[str, Dict[str, Any]]: A dictionary with table and context information.
    """
    db_info: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = {}

    def populate_metadata(md: MetaData) -> None:
        # Create empty metadata dictionary with column datatypes
        for table_name, table in md.tables.items():
            # Populate tables with empty descriptions
            db_info[table_name] = {
                "description": info[table_name]["description"] or "",
                "columns": {},
            }

            for column in table.columns:
                # Populate columns with datatype
                db_info[table_name]["columns"][str(column.name)] = (  # type: ignore
                    str(column.type)
                )

    if isinstance(metadata, list):
        for md in metadata:
            populate_metadata(md)
    else:
        populate_metadata(metadata)

    return db_info


def populate_metadata(
    metadata: MetaData | List[MetaData],
    info: Dict[str, Dict[str, Union[str, Dict[str, str]]]],
) -> Dict[str, Dict[str, Union[str, Dict[str, str]]]]:
    """
    Populate metadata based on the provided database metadata and additional info.

    Args:
        metadata (MetaData): Metadata object from SQLAlchemy.
        info (Dict): Additional information for database tables and columns.

    Returns:
        Dict: A dictionary containing populated metadata information.
    """
    # Fetch basic metadata info using available tools
    db_info: Dict[str, Dict[str, Union[str, Dict[str, str]]]] = (
        populate_metadata_with_schema_tools(metadata=metadata, info=info)
    )

    # Iterate over tables to update column metadata
    for table_name in db_info.keys():
        # Update only if additional info for the table exists
        if table_name in info:
            for column_name in db_info[table_name]["columns"]:
                # Merge and update column description if available
                if column_name in info[table_name]["columns"]:
                    db_info[table_name]["columns"][column_name] = (  # type: ignore
                        db_info[table_name]["columns"][column_name]  # type: ignore
                        + "; "
                        + info[table_name]["columns"][column_name]  # type: ignore
                    )

    return db_info
</file>

<file path="langroid/agent/special/sql/utils/system_message.py">
DEFAULT_SYS_MSG = """You are a savvy data scientist/database administrator, 
with expertise in answering questions by querying a {dialect} database.
You do not have access to the database 'db' directly, so you will need to use the 
`run_query` tool/function-call to answer questions.

The below JSON schema maps the SQL database structure. It outlines tables, each 
with a description and columns. Each table is identified by a key, 
and holds a description and a dictionary of columns, 
with column names as keys and their descriptions as values. 
{schema_dict}

ONLY the tables and column names and tables specified above should be used in
the generated queries. 
You must be smart about using the right tables and columns based on the 
english description. If you are thinking of using a table or column that 
does not exist, you are probably on the wrong track, so you should try
your best to answer based on an existing table or column.
DO NOT assume any tables or columns other than those above."""

SCHEMA_TOOLS_SYS_MSG = """You are a savvy data scientist/database administrator, 
with expertise in answering questions by interacting with a SQL database.

You will have to follow these steps to complete your job:
1) Use the `get_table_names` tool/function-call to get a list of all possibly 
relevant table names.
2) Use the `get_table_schema` tool/function-call to get the schema of all 
possibly relevant tables to identify possibly relevant columns. Only 
call this method on potentially relevant tables.
3) Use the `get_column_descriptions` tool/function-call to get more information 
about any relevant columns.
4) Write a {dialect} query and use `run_query` tool the Execute the SQL query 
on the database to obtain the results.

Do not make assumptions about the database schema before using the tools.
Use the tool/functions to learn more about the database schema."""
</file>

<file path="langroid/agent/special/sql/utils/tools.py">
from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage


class RunQueryTool(ToolMessage):
    request: str = "run_query"
    purpose: str = """
            To run <query> on the database 'db' and 
            return the results to answer a question.
            """
    query: str

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            cls(
                query="SELECT * FROM movies WHERE genre = 'comedy'",
            ),
            (
                "Find all movies with a rating of 5",
                cls(
                    query="SELECT * FROM movies WHERE rating = 5",
                ),
            ),
        ]


class GetTableNamesTool(ToolMessage):
    request: str = "get_table_names"
    purpose: str = """
            To retrieve the names of all <tables> in the database 'db'.
            """


class GetTableSchemaTool(ToolMessage):
    request: str = "get_table_schema"
    purpose: str = """
            To retrieve the schema of all provided <tables> in the database 'db'.
            """
    tables: List[str]

    @classmethod
    def example(cls) -> "GetTableSchemaTool":
        return cls(
            tables=["employees", "departments", "sales"],
        )


class GetColumnDescriptionsTool(ToolMessage):
    request: str = "get_column_descriptions"
    purpose: str = """
            To retrieve the description of one or more <columns> from the respective 
            <table> in the database 'db'.
            """
    table: str
    columns: str

    @classmethod
    def example(cls) -> "GetColumnDescriptionsTool":
        return cls(
            table="employees",
            columns="name, department_id",
        )
</file>

<file path="langroid/agent/special/sql/__init__.py">
from . import utils


__all__ = [
    "utils",
]

try:
    from . import sql_chat_agent
    from .sql_chat_agent import SQLChatAgentConfig, SQLChatAgent

    sql_chat_agent
    SQLChatAgent
    SQLChatAgentConfig
    __all__.extend(["SQLChatAgentConfig", "SQLChatAgent", "sql_chat_agent"])
except ImportError:
    pass
</file>

<file path="langroid/agent/special/__init__.py">
from .relevance_extractor_agent import (
    RelevanceExtractorAgent,
    RelevanceExtractorAgentConfig,
)
from .doc_chat_agent import DocChatAgent, DocChatAgentConfig
from .retriever_agent import (
    RecordMetadata,
    RecordDoc,
    RetrieverAgentConfig,
    RetrieverAgent,
)
from .lance_doc_chat_agent import LanceDocChatAgent
from .table_chat_agent import (
    dataframe_summary,
    TableChatAgent,
    TableChatAgentConfig,
    PandasEvalTool,
)


from . import relevance_extractor_agent
from . import doc_chat_agent
from . import retriever_agent
from . import lance_tools
from . import lance_doc_chat_agent
from . import lance_rag
from . import table_chat_agent


__all__ = [
    "RelevanceExtractorAgent",
    "RelevanceExtractorAgentConfig",
    "DocChatAgent",
    "DocChatAgentConfig",
    "RecordMetadata",
    "RecordDoc",
    "RetrieverAgentConfig",
    "RetrieverAgent",
    "dataframe_summary",
    "TableChatAgent",
    "TableChatAgentConfig",
    "PandasEvalTool",
    "relevance_extractor_agent",
    "doc_chat_agent",
    "retriever_agent",
    "table_chat_agent",
    "LanceDocChatAgent",
    "lance_tools",
    "lance_doc_chat_agent",
    "lance_rag",
]

try:
    from . import sql

    sql
    __all__.append("sql")
except ImportError:
    pass
</file>

<file path="langroid/agent/special/lance_tools.py">
import logging

from pydantic import BaseModel, Field

from langroid.agent.tool_message import ToolMessage

logger = logging.getLogger(__name__)


class QueryPlan(BaseModel):
    original_query: str = Field(..., description="The original query for reference")
    query: str = Field(..., description="A possibly NON-EMPTY rephrased query")
    filter: str = Field(
        "",
        description="Filter condition if needed (or empty if no filter is needed)",
    )
    dataframe_calc: str = Field(
        "", description="An optional Pandas-dataframe calculation/aggregation string"
    )


class QueryPlanTool(ToolMessage):
    request: str = "query_plan"  # the agent method name that handles this tool
    purpose: str = """
    Given a user's query, generate a query <plan> consisting of:
    - <original_query> - the original query for reference
    - <filter> condition if needed (or empty string if no filter is needed)
    - <query> - a possibly NON-EMPTY rephrased query that can be used to match the 
        CONTENT of the documents 
        (can be same as <original_query> if no rephrasing is needed)
    - <dataframe_calc> - a Pandas-dataframe calculation/aggregation string
        that can be used to calculate the answer 
        (or empty string if no calculation is needed).
    """
    plan: QueryPlan


class AnswerTool(ToolMessage):
    """Wrapper for answer from LanceDocChatAgent"""

    purpose: str = "To package the answer from LanceDocChatAgent"
    request: str = "answer_tool"
    answer: str


class QueryPlanAnswerTool(ToolMessage):
    request: str = "query_plan_answer"  # the agent method name that handles this tool
    purpose: str = """
    Assemble query <plan> and <answer>
    """
    plan: QueryPlan
    answer: str = Field(..., description="The answer received from the assistant")


class QueryPlanFeedbackTool(ToolMessage):
    request: str = "query_plan_feedback"
    purpose: str = """
    To give <feedback> regarding the query plan, 
    along with a <suggested_fix> if any (empty string if no fix is suggested).
    """
    feedback: str
    suggested_fix: str
</file>

<file path="langroid/agent/special/relevance_extractor_agent.py">
"""
Agent to retrieve relevant segments from a body of text,
that are relevant to a query.

"""

import logging
from typing import Optional, no_type_check

from rich.console import Console

from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.chat_document import ChatDocument
from langroid.agent.tools.segment_extract_tool import SegmentExtractTool
from langroid.language_models.base import LLMConfig
from langroid.language_models.openai_gpt import OpenAIGPTConfig
from langroid.parsing.utils import extract_numbered_segments, number_segments
from langroid.utils.constants import DONE, NO_ANSWER

console = Console()
logger = logging.getLogger(__name__)


class RelevanceExtractorAgentConfig(ChatAgentConfig):
    llm: LLMConfig | None = OpenAIGPTConfig()
    segment_length: int = 1  # number of sentences per segment
    query: str = ""  # query for relevance extraction
    handle_llm_no_tool: str = """
    You FORGOT to use the `extract_segments` tool!
    Remember that your response MUST be a JSON-formatted string
    starting with `{"request": "extract_segments", ...}`
    """
    system_message: str = """
    The user will give you a PASSAGE containing segments numbered as  
    <#1#>, <#2#>, <#3#>, etc.,
    followed by a QUERY. Extract ONLY the segment-numbers from 
    the PASSAGE that are RELEVANT to the QUERY.
    Present the extracted segment-numbers using the `extract_segments` tool/function.
    Note that your response MUST be a JSON-formatted string 
    starting with `{"request": "extract_segments", ...}`
    """


class RelevanceExtractorAgent(ChatAgent):
    """
    Agent for extracting segments from text, that are relevant to a given query.
    """

    def __init__(self, config: RelevanceExtractorAgentConfig):
        super().__init__(config)
        self.config: RelevanceExtractorAgentConfig = config
        self.enable_message(SegmentExtractTool)
        self.numbered_passage: Optional[str] = None

    @no_type_check
    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        """Compose a prompt asking to extract relevant segments from a passage.
        Steps:
        - number the segments in the passage
        - compose prompt
        - send to LLM
        """
        assert self.config.query is not None, "No query specified"
        assert message is not None, "No message specified"
        message_str = message.content if isinstance(message, ChatDocument) else message
        # number the segments in the passage
        self.numbered_passage = number_segments(message_str, self.config.segment_length)
        # compose prompt
        prompt = f"""
        <Instructions>
        Given the PASSAGE below with NUMBERED segments, and the QUERY,
        extract ONLY the segment-numbers that are RELEVANT to the QUERY,
        and present them using the `extract_segments` tool/function,
        i.e. your response MUST be a JSON-formatted string starting with
        `{{"request": "extract_segments", ...}}`
        </Instructions>
        
        PASSAGE:
        {self.numbered_passage}
        
        QUERY: {self.config.query}
        """
        # send to LLM
        response = super().llm_response(prompt)
        return response

    @no_type_check
    async def llm_response_async(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        """
        Compose a prompt asking to extract relevant segments from a passage.
        Steps:
        - number the segments in the passage
        - compose prompt
        - send to LLM
        The LLM is expected to generate a structured msg according to the
        SegmentExtractTool schema, i.e. it should contain a `segment_list` field
        whose value is a list of segment numbers or ranges, like "10,12,14-17".
        """

        assert self.config.query is not None, "No query specified"
        assert message is not None, "No message specified"
        message_str = message.content if isinstance(message, ChatDocument) else message
        # number the segments in the passage
        self.numbered_passage = number_segments(message_str, self.config.segment_length)
        # compose prompt
        prompt = f"""
        PASSAGE:
        {self.numbered_passage}
        
        QUERY: {self.config.query}
        """
        # send to LLM
        response = await super().llm_response_async(prompt)
        return response

    def extract_segments(self, msg: SegmentExtractTool) -> str:
        """Method to handle a segmentExtractTool message from LLM"""
        spec = msg.segment_list
        if len(self.message_history) == 0:
            return DONE + " " + NO_ANSWER
        if spec is None or spec.strip() in ["", NO_ANSWER]:
            return DONE + " " + NO_ANSWER
        assert self.numbered_passage is not None, "No numbered passage"
        # assume this has numbered segments
        try:
            extracts = extract_numbered_segments(self.numbered_passage, spec)
        except Exception:
            return DONE + " " + NO_ANSWER
        # this response ends the task by saying DONE
        return DONE + " " + extracts
</file>

<file path="langroid/agent/special/retriever_agent.py">
"""
DEPRECATED: use DocChatAgent instead, with DocChatAgentConfig.retrieve_only=True,
and if you want to retrieve FULL relevant doc-contents rather than just extracts,
then set DocChatAgentConfig.extraction_granularity=-1

This is an agent to retrieve relevant extracts from a vector store,
where the LLM is used to filter for "true" relevance after retrieval from the
vector store.
This is essentially the same as DocChatAgent, except that instead of
generating final summary answer based on relevant extracts, it just returns
those extracts.
See test_retriever_agent.py for example usage.
"""

import logging
from typing import Sequence

from rich.console import Console

from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
from langroid.mytypes import DocMetaData, Document

console = Console()
logger = logging.getLogger(__name__)

# for backwards compatibility:
RecordMetadata = DocMetaData
RecordDoc = Document
RetrieverAgentConfig = DocChatAgentConfig


class RetrieverAgent(DocChatAgent):
    """
    Agent for just retrieving chunks/docs/extracts matching a query
    """

    def __init__(self, config: DocChatAgentConfig):
        super().__init__(config)
        self.config: DocChatAgentConfig = config
        logger.warning(
            """
        `RetrieverAgent` is deprecated. Use `DocChatAgent` instead, with
        `DocChatAgentConfig.retrieve_only=True`, and if you want to retrieve
        FULL relevant doc-contents rather than just extracts, then set
        `DocChatAgentConfig.extraction_granularity=-1`
        """
        )

    def get_records(self) -> Sequence[Document]:
        # subclasses should override
        return []

    def ingest(self) -> None:
        records = self.get_records()
        if self.vecdb is None:
            logger.warning("Vector store not configured. Cannot ingest records.")
        else:
            self.vecdb.add_documents(records)
</file>

<file path="langroid/agent/tools/mcp/__init__.py">
from .decorators import mcp_tool
from .fastmcp_client import (
    FastMCPClient,
    get_tool,
    get_tool_async,
    get_tools,
    get_tools_async,
    get_mcp_tool_async,
    get_mcp_tools_async,
)


__all__ = [
    "mcp_tool",
    "FastMCPClient",
    "get_tool",
    "get_tool_async",
    "get_tools",
    "get_tools_async",
    "get_mcp_tool_async",
    "get_mcp_tools_async",
]
</file>

<file path="langroid/agent/tools/mcp/decorators.py">
from typing import Callable, Type

from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.mcp.fastmcp_client import (
    FastMCPServerSpec,
    get_tool,
)


def mcp_tool(
    server: FastMCPServerSpec, tool_name: str
) -> Callable[[Type[ToolMessage]], Type[ToolMessage]]:
    """Decorator: declare a ToolMessage class bound to a FastMCP tool.

    Usage:
        @mcp_tool("/path/to/server.py", "get_weather")
        class WeatherTool:
            def pretty(self) -> str:
                return f"Temp is {self.temperature}"

    The `server` may be a string/URL/FastMCP/ClientTransport, or a zero-arg
    callable returning one of those, e.g. `lambda: StdioTransport(...)`. Using a
    factory ensures a fresh transport per connection under fastmcp>=2.13.
    """

    def decorator(user_cls: Type[ToolMessage]) -> Type[ToolMessage]:
        # build the “real” ToolMessage subclass for this server/tool
        RealTool: Type[ToolMessage] = get_tool(server, tool_name)

        # copy user‐defined methods / attributes onto RealTool
        for name, attr in user_cls.__dict__.items():
            if name.startswith("__") and name.endswith("__"):
                continue
            setattr(RealTool, name, attr)

        # preserve the user’s original name if you like:
        RealTool.__name__ = user_cls.__name__
        return RealTool

    return decorator
</file>

<file path="langroid/agent/tools/__init__.py">
from . import google_search_tool
from . import recipient_tool
from . import rewind_tool
from . import orchestration
from .google_search_tool import GoogleSearchTool
from .recipient_tool import AddRecipientTool, RecipientTool
from .rewind_tool import RewindTool
from .orchestration import (
    AgentDoneTool,
    DoneTool,
    ForwardTool,
    PassTool,
    SendTool,
    AgentSendTool,
    DonePassTool,
    ResultTool,
    FinalResultTool,
)

__all__ = [
    "GoogleSearchTool",
    "AddRecipientTool",
    "RecipientTool",
    "google_search_tool",
    "recipient_tool",
    "rewind_tool",
    "RewindTool",
    "orchestration",
    "AgentDoneTool",
    "DoneTool",
    "DonePassTool",
    "ForwardTool",
    "PassTool",
    "SendTool",
    "AgentSendTool",
    "ResultTool",
    "FinalResultTool",
]
</file>

<file path="langroid/agent/tools/duckduckgo_search_tool.py">
"""
A tool to trigger a DuckDuckGo search for a given query, and return the top results with
their titles, links, summaries. Since the tool is stateless (i.e. does not need
access to agent state), it can be enabled for any agent, without having to define a
special method inside the agent: `agent.enable_message(DuckduckgoSearchTool)`
"""

from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage
from langroid.parsing.web_search import duckduckgo_search


class DuckduckgoSearchTool(ToolMessage):
    request: str = "duckduckgo_search"
    purpose: str = """
            To search the web and return up to <num_results> 
            links relevant to the given <query>. When using this tool,
            ONLY show the required JSON, DO NOT SAY ANYTHING ELSE.
            Wait for the results of the web search, and then use them to
            compose your response.
            """
    query: str
    num_results: int

    def handle(self) -> str:
        """
        Conducts a search using DuckDuckGo based on the provided query
        and number of results by triggering a duckduckgo_search.

        Returns:
            str: A formatted string containing the titles, links, and
                summaries of each search result, separated by two newlines.
        """
        search_results = duckduckgo_search(self.query, self.num_results)
        # return Title, Link, Summary of each result, separated by two newlines
        results_str = "\n\n".join(str(result) for result in search_results)
        return f"""
        BELOW ARE THE RESULTS FROM THE WEB SEARCH. USE THESE TO COMPOSE YOUR RESPONSE:
        {results_str}
        """

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            cls(
                query="When was the Llama2 Large Language Model (LLM) released?",
                num_results=3,
            ),
        ]
</file>

<file path="langroid/agent/tools/exa_search_tool.py">
"""
A tool to trigger a Exa search for a given query,
(https://docs.exa.ai/reference/getting-started)
and return the top results with their titles, links, summaries.
Since the tool is stateless (i.e. does not need
access to agent state), it can be enabled for any agent, without having to define a
special method inside the agent: `agent.enable_message(ExaSearchTool)`

NOTE: To use this tool, you need to:

* set the EXA_API_KEY environment variables in
your `.env` file, e.g. `EXA_API_KEY=your_api_key_here`
(Note as of 28 Jan 2023, Metaphor renamed to Exa, so you can also use
`EXA_API_KEY=your_api_key_here`)

* install langroid with the `exa-py` extra, e.g.
`pip install langroid[exa]` or `uv pip install langroid[exa]`
or `poetry add langroid[exa]`  or `uv add langroid[exa]`
(it installs the `exa_py` package from pypi).

For more information, please refer to the official docs:
https://exa.ai/
"""

from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage
from langroid.parsing.web_search import exa_search


class ExaSearchTool(ToolMessage):
    request: str = "exa_search"
    purpose: str = """
            To search the web and return up to <num_results> 
            links relevant to the given <query>. When using this tool,
            ONLY show the required JSON, DO NOT SAY ANYTHING ELSE.
            Wait for the results of the web search, and then use them to
            compose your response.
            """
    query: str
    num_results: int

    def handle(self) -> str:
        """
        Conducts a search using the exa API based on the provided query
        and number of results by triggering a exa_search.

        Returns:
            str: A formatted string containing the titles, links, and
                summaries of each search result, separated by two newlines.
        """

        search_results = exa_search(self.query, self.num_results)
        # return Title, Link, Summary of each result, separated by two newlines
        results_str = "\n\n".join(str(result) for result in search_results)
        return f"""
        BELOW ARE THE RESULTS FROM THE WEB SEARCH. USE THESE TO COMPOSE YOUR RESPONSE:
        {results_str}
        """

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            cls(
                query="When was the Llama2 Large Language Model (LLM) released?",
                num_results=3,
            ),
        ]
</file>

<file path="langroid/agent/tools/google_search_tool.py">
"""
A tool to trigger a Google search for a given query, and return the top results with
their titles, links, summaries. Since the tool is stateless (i.e. does not need
access to agent state), it can be enabled for any agent, without having to define a
special method inside the agent: `agent.enable_message(GoogleSearchTool)`

NOTE: Using this tool requires setting the GOOGLE_API_KEY and GOOGLE_CSE_ID
environment variables in your `.env` file, as explained in the
[README](https://github.com/langroid/langroid#gear-installation-and-setup).
"""

from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage
from langroid.parsing.web_search import google_search


class GoogleSearchTool(ToolMessage):
    request: str = "web_search"
    purpose: str = """
            To search the web and return up to <num_results> links relevant to 
            the given <query>. 
            """
    query: str
    num_results: int

    def handle(self) -> str:
        search_results = google_search(self.query, self.num_results)
        # return Title, Link, Summary of each result, separated by two newlines
        return "\n\n".join(str(result) for result in search_results)

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            cls(
                query="When was the Llama2 Large Language Model (LLM) released?",
                num_results=3,
            ),
        ]
</file>

<file path="langroid/agent/tools/retrieval_tool.py">
from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage


class RetrievalTool(ToolMessage):
    """
    Retrieval tool, only to be used by a DocChatAgent.
    The handler method is defined in DocChatAgent.retrieval_tool
    """

    request: str = "retrieval_tool"
    purpose: str = """
            To retrieve up to <num_results> passages from a document-set, that are 
            relevant to a <query>, which could be a question or simply a topic or 
            search phrase. 
            """
    query: str
    num_results: int

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            cls(
                query="What are the eligibility criteria for the scholarship?",
                num_results=3,
            ),
            cls(
                query="Self-Attention mechanism in RNNs",
                num_results=5,
            ),
        ]
</file>

<file path="langroid/agent/tools/segment_extract_tool.py">
"""
A tool to extract segment numbers from the last user message,
containing numbered segments.

The idea is that when an LLM wants to (or is asked to) simply extract
portions of a message verbatim, it should use this tool/function to
SPECIFY what should be extracted, rather than actually extracting it.
The output will be in the form of a list of segment numbers or ranges.
This will usually be much cheaper and faster than actually writing out the extracted
text. The handler of this tool/function will then extract the text and send it back.
"""

from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage


class SegmentExtractTool(ToolMessage):
    request: str = "extract_segments"
    purpose: str = """
            To extract segments from a body of text containing numbered 
            segments, in the form of a <segment_list> which is a list of segment 
            numbers or ranges, like "10,12,14-17".
            """
    segment_list: str

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            (
                "I want to extract segments 1, 3, and 5 thru 7",
                cls(segment_list="1,3,5-7"),
            )
        ]

    @classmethod
    def instructions(cls) -> str:
        return """
        Use this tool/function to indicate certain segments from 
        a body of text containing numbered segments.
        """
</file>

<file path="langroid/agent/tools/task_tool.py">
"""
TaskTool: A tool that allows agents to delegate a task to a sub-agent with
    specific tools enabled.
"""

import uuid
from typing import List, Optional

from pydantic import Field
from pydantic.fields import ModelPrivateAttr

import langroid.language_models as lm
from langroid import ChatDocument
from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
from langroid.agent.task import Task
from langroid.agent.tool_message import ToolMessage
from langroid.agent.tools.orchestration import DoneTool


class TaskTool(ToolMessage):
    """
    Tool that spawns a sub-agent with specified tools to handle a task.

    The sub-agent can be given a custom name for identification in logs.
    If no name is provided, a random unique name starting with 'agent'
    will be generated.
    """

    # TODO: setting up termination conditions of sub-task needs to be improved
    request: str = "task_tool"
    purpose: str = """
        <HowToUse>
        Use this tool to delegate a task to a sub-agent with specific tools enabled.
        The sub-agent will be created with the specified tools and will run the task
        non-interactively.
    """

    # Parameters for the agent tool

    system_message: Optional[str] = Field(
        ...,
        description="""
        Optional system message to configure the sub-agent's general behavior and 
        to specify the task and its context.
            A good system message will have these components:
            - Inform the sub-agent of its role, e.g. "You are a financial analyst."
            - Clear spec of the task, with sufficient context for the sub-agent to 
              understand what it needs to do, since the sub-agent does 
              NOT have access to your conversation history!
            - Any additional general context needed for the task, such as a
              (part of a) document, or data items, etc.
            - Specify when to use certain tools, e.g. 
                "You MUST use the 'stock_data' tool to extract stock information.
        """,
    )

    prompt: str = Field(
        ...,
        description="""
            The prompt to run the sub-agent with. This differs from the agent's
            system message: Whereas the system message configures the sub-agent's
            GENERAL role and goals, the `prompt` is the SPECIFIC input that the 
            sub-agent will process. In LLM terms, the system message is sent to the 
            LLM as the first message, with role = "system" or "developer", and 
            the prompt is sent as a message with role = "user".
            EXAMPLE: system_message = "You are a financial analyst, when the 
                user asks about the share-price of a company, 
                you must use your tools to do the research, and 
                return the final answer to the user."
            
            prompt = "What is the share-price of Apple Inc.?"
            """,
    )

    tools: List[str] = Field(
        ...,
        description="""
        A list of tool names to enable for the sub-agent.
        This must be a list of strings referring to the names of tools
        that are known to you. 
        If you want to enable all tools, or you do not have any preference
        on what tools are enabled for the sub-agent, you can set 
        this field to a singleton list ['ALL']
        To disable all tools, set it to a singleton list ['NONE']
        """,
    )
    # TODO: ensure valid model name
    model: Optional[str] = Field(
        default=None,
        description="""
            Optional name of the LLM model to use for the sub-agent, e.g. 'gpt-4.1'
            If omitted, the sub-agent will use the same model as yours.
            """,
    )
    max_iterations: Optional[int] = Field(
        default=None,
        description="Optional max iterations for the sub-agent to run the task",
    )
    agent_name: Optional[str] = Field(
        default=None,
        description="""
            Optional name for the sub-agent. This will be used as the agent's name
            in logs and for identification purposes. If not provided, a random unique
            name starting with 'agent' will be generated.
            """,
    )

    def _set_up_task(self, agent: ChatAgent) -> Task:
        """
        Helper method to set up a task for the sub-agent.

        Args:
            agent: The parent ChatAgent that is handling this tool
        """
        # Generate a random name if not provided
        agent_name = self.agent_name or f"agent-{str(uuid.uuid4())[:8]}"

        # Create chat agent config with system message if provided
        # TODO: Maybe we just copy the parent agent's config and override chat_model?
        #   -- but what if parent agent has a MockLMConfig?
        llm_config = lm.OpenAIGPTConfig(
            chat_model=self.model or lm.OpenAIChatModel.GPT4_1_MINI,
        )
        config = ChatAgentConfig(
            name=agent_name,
            llm=llm_config,
            handle_llm_no_tool=f"""
                You forgot to use one of your TOOLs! Remember that you must either:
                - use a tool, or a sequence of tools, to complete your task, OR
                - if you are done with your task, use the `{DoneTool.name()}` tool
                to return the result.
                
                As a reminder, this was your task:
                {self.prompt}
                """,
            system_message=f"""
                {self.system_message}
                
                When you are finished with your task, you MUST
                use the TOOL `{DoneTool.name()}` to end the task
                and return the result.                
            """,
        )

        # Create the sub-agent
        sub_agent = ChatAgent(config)

        # Enable the specified tools for the sub-agent
        # Convert tool names to actual tool classes using parent agent's tools_map
        if self.tools == ["ALL"]:
            # Enable all tools from the parent agent:
            # This is the list of all tools KNOWN (whether usable or handle-able or not)
            tool_classes = []
            for t in agent.llm_tools_known:
                if t in agent.llm_tools_map and t != self.request:
                    tool_class = agent.llm_tools_map[t]
                    allow_llm_use = tool_class._allow_llm_use
                    if isinstance(allow_llm_use, ModelPrivateAttr):
                        allow_llm_use = allow_llm_use.default
                    if allow_llm_use:
                        tool_classes.append(tool_class)
        elif self.tools == ["NONE"]:
            # No tools enabled
            tool_classes = []
        else:
            # Enable only specified tools
            tool_classes = []
            for tool_name in self.tools:
                if tool_name in agent.llm_tools_map:
                    tool_class = agent.llm_tools_map[tool_name]
                    allow_llm_use = tool_class._allow_llm_use
                    if isinstance(allow_llm_use, ModelPrivateAttr):
                        allow_llm_use = allow_llm_use.default
                    if allow_llm_use:
                        tool_classes.append(tool_class)

        # always enable the DoneTool to signal task completion
        sub_agent.enable_message(tool_classes + [DoneTool], use=True, handle=True)

        # Create a non-interactive task
        task = Task(sub_agent, interactive=False)

        return task

    def handle(
        self, agent: ChatAgent, chat_doc: Optional[ChatDocument] = None
    ) -> Optional[ChatDocument]:
        """

        Handle the TaskTool by creating a sub-agent with specified tools
        and running the task non-interactively.

        Args:
            agent: The parent ChatAgent that is handling this tool
            chat_doc: The ChatDocument containing this tool message
        """

        task = self._set_up_task(agent)

        # Create a ChatDocument for the prompt with parent pointer
        prompt_doc = None
        if chat_doc is not None:
            from langroid.agent.chat_document import ChatDocMetaData

            prompt_doc = ChatDocument(
                content=self.prompt,
                metadata=ChatDocMetaData(
                    parent_id=chat_doc.id(),
                    agent_id=agent.id,
                    sender=chat_doc.metadata.sender,
                ),
            )
            # Set bidirectional parent-child relationship
            chat_doc.metadata.child_id = prompt_doc.id()

        # Run the task with the ChatDocument or string prompt
        result = task.run(prompt_doc or self.prompt, turns=self.max_iterations or 10)
        return result

    async def handle_async(
        self, agent: ChatAgent, chat_doc: Optional[ChatDocument] = None
    ) -> Optional[ChatDocument]:
        """
        Async method to handle the TaskTool by creating a sub-agent with specified tools
        and running the task non-interactively.

        Args:
            agent: The parent ChatAgent that is handling this tool
            chat_doc: The ChatDocument containing this tool message
        """
        task = self._set_up_task(agent)

        # Create a ChatDocument for the prompt with parent pointer
        prompt_doc = None
        if chat_doc is not None:
            from langroid.agent.chat_document import ChatDocMetaData

            prompt_doc = ChatDocument(
                content=self.prompt,
                metadata=ChatDocMetaData(
                    parent_id=chat_doc.id(),
                    agent_id=agent.id,
                    sender=chat_doc.metadata.sender,
                ),
            )
            # Set bidirectional parent-child relationship
            chat_doc.metadata.child_id = prompt_doc.id()

        # Run the task with the ChatDocument or string prompt
        # TODO eventually allow the various task setup configs,
        #  including termination conditions
        result = await task.run_async(
            prompt_doc or self.prompt, turns=self.max_iterations or 10
        )
        return result
</file>

<file path="langroid/agent/tools/tavily_search_tool.py">
"""
A tool to trigger a Tavily search for a given query, and return the top results with
their titles, links, summaries. Since the tool is stateless (i.e. does not need
access to agent state), it can be enabled for any agent, without having to define a
special method inside the agent: `agent.enable_message(TavilySearchTool)`
"""

from typing import List, Tuple

from langroid.agent.tool_message import ToolMessage
from langroid.parsing.web_search import tavily_search


class TavilySearchTool(ToolMessage):
    request: str = "tavily_search"
    purpose: str = """
            To search the web and return up to <num_results> 
            links relevant to the given <query>. When using this tool,
            ONLY show the required JSON, DO NOT SAY ANYTHING ELSE.
            Wait for the results of the web search, and then use them to
            compose your response.
            """
    query: str
    num_results: int

    def handle(self) -> str:
        """
        Conducts a search using Tavily based on the provided query
        and number of results by triggering a tavily_search.

        Returns:
            str: A formatted string containing the titles, links, and
                summaries of each search result, separated by two newlines.
        """
        search_results = tavily_search(self.query, self.num_results)
        # return Title, Link, Summary of each result, separated by two newlines
        results_str = "\n\n".join(str(result) for result in search_results)
        return f"""
        BELOW ARE THE RESULTS FROM THE WEB SEARCH. USE THESE TO COMPOSE YOUR RESPONSE:
        {results_str}
        """

    @classmethod
    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
        return [
            cls(
                query="When was the Llama2 Large Language Model (LLM) released?",
                num_results=3,
            ),
        ]
</file>

<file path="langroid/agent/__init__.py">
from .base import Agent, AgentConfig
from .chat_document import (
    ChatDocAttachment,
    ChatDocMetaData,
    ChatDocLoggerFields,
    ChatDocument,
)
from .chat_agent import ChatAgentConfig, ChatAgent
from .tool_message import ToolMessage
from .task import Task

from . import base
from . import chat_document
from . import chat_agent
from . import task
from . import batch
from . import tool_message
from . import tools
from . import special


__all__ = [
    "Agent",
    "AgentConfig",
    "ChatDocAttachment",
    "ChatDocMetaData",
    "ChatDocLoggerFields",
    "ChatDocument",
    "ChatAgent",
    "ChatAgentConfig",
    "ToolMessage",
    "Task",
    "base",
    "chat_document",
    "chat_agent",
    "task",
    "batch",
    "tool_message",
    "tools",
    "special",
]
</file>

<file path="langroid/agent/done_sequence_parser.py">
"""Parser for done sequence DSL (Domain Specific Language).

Converts string patterns into DoneSequence objects for convenient task completion
configuration.

Examples:
    "T, A" -> Tool followed by Agent response
    "T[calculator], A" -> Specific tool 'calculator' followed by Agent response
    "L, T, A, L" -> LLM, Tool, Agent, LLM sequence
    "C[quit|exit]" -> Content matching regex pattern
"""

import re
from typing import Any, Dict, List, Optional, Union

from .task import AgentEvent, DoneSequence, EventType


def parse_done_sequence(
    sequence: Union[str, DoneSequence], tools_map: Optional[Dict[str, Any]] = None
) -> DoneSequence:
    """Parse a string pattern or return existing DoneSequence unchanged.

    Args:
        sequence: Either a DoneSequence object or a string pattern to parse
        tools_map: Optional dict mapping tool names to tool classes
            (e.g., agent.llm_tools_map)

    Returns:
        DoneSequence object

    Raises:
        ValueError: If the string pattern is invalid
    """
    if isinstance(sequence, DoneSequence):
        return sequence

    if not isinstance(sequence, str):
        raise ValueError(f"Expected string or DoneSequence, got {type(sequence)}")

    events = _parse_string_pattern(sequence, tools_map)
    return DoneSequence(events=events)


def _parse_string_pattern(
    pattern: str, tools_map: Optional[Dict[str, Any]] = None
) -> List[AgentEvent]:
    """Parse a string pattern into a list of AgentEvent objects.

    Pattern format:
        - Single letter codes: T, A, L, U, N, C
        - Specific tools: T[tool_name] or T[ToolClass]
        - Content match: C[regex_pattern]
        - Separated by commas, spaces allowed

    Args:
        pattern: String pattern to parse
        tools_map: Optional dict mapping tool names to tool classes
            (e.g., agent.llm_tools_map)

    Returns:
        List of AgentEvent objects

    Raises:
        ValueError: If pattern is invalid
    """
    events = []

    # Split by comma and strip whitespace
    parts = [p.strip() for p in pattern.split(",")]

    for part in parts:
        if not part:
            continue

        event = _parse_event_token(part, tools_map)
        events.append(event)

    if not events:
        raise ValueError(f"No valid events found in pattern: {pattern}")

    return events


def _parse_event_token(
    token: str, tools_map: Optional[Dict[str, Any]] = None
) -> AgentEvent:
    """Parse a single event token into an AgentEvent.

    Args:
        token: Single event token (e.g., "T", "T[calc]", "C[quit|exit]")
        tools_map: Optional dict mapping tool names to tool classes
            (e.g., agent.llm_tools_map)

    Returns:
        AgentEvent object

    Raises:
        ValueError: If token is invalid
    """
    # Check for bracket notation
    bracket_match = re.match(r"^([A-Z])\[([^\]]+)\]$", token)

    if bracket_match:
        event_code = bracket_match.group(1)
        param = bracket_match.group(2)

        if event_code == "T":
            # Specific tool: T[tool_name] or T[ToolClass]
            tool_class = None
            tool_name = param

            # First try direct lookup in tools_map by the param (tool name)
            if tools_map and param in tools_map:
                tool_class = tools_map[param]
                tool_name = param
            elif tools_map:
                # If not found, loop through tools_map to find a tool class
                # whose __name__ matches param
                for name, cls in tools_map.items():
                    if hasattr(cls, "__name__") and cls.__name__ == param:
                        tool_class = cls
                        tool_name = name
                        break

            return AgentEvent(
                event_type=EventType.SPECIFIC_TOOL,
                tool_name=tool_name,
                tool_class=tool_class,
            )
        elif event_code == "C":
            # Content match: C[regex_pattern]
            return AgentEvent(event_type=EventType.CONTENT_MATCH, content_pattern=param)
        else:
            raise ValueError(
                f"Invalid event code with brackets: {event_code}. "
                "Only T[tool] and C[pattern] are supported."
            )

    # Simple single-letter codes
    event_map = {
        "T": EventType.TOOL,
        "A": EventType.AGENT_RESPONSE,
        "L": EventType.LLM_RESPONSE,
        "U": EventType.USER_RESPONSE,
        "N": EventType.NO_RESPONSE,
        "C": EventType.CONTENT_MATCH,  # C without brackets matches any content
    }

    if token in event_map:
        return AgentEvent(event_type=event_map[token])

    # If not a single letter, could be a full event type name
    token_upper = token.upper()
    if token_upper == "TOOL":
        return AgentEvent(event_type=EventType.TOOL)
    elif token_upper == "AGENT":
        return AgentEvent(event_type=EventType.AGENT_RESPONSE)
    elif token_upper == "LLM":
        return AgentEvent(event_type=EventType.LLM_RESPONSE)
    elif token_upper == "USER":
        return AgentEvent(event_type=EventType.USER_RESPONSE)
    else:
        raise ValueError(
            f"Invalid event token: '{token}'. "
            "Valid tokens are: T, A, L, U, N, C, or T[tool_name], C[pattern]"
        )


def parse_done_sequences(
    sequences: List[Union[str, DoneSequence]],
    tools_map: Optional[Dict[str, Any]] = None,
) -> List[DoneSequence]:
    """Parse a list of mixed string patterns and DoneSequence objects.

    Args:
        sequences: List containing strings and/or DoneSequence objects
        tools_map: Optional dict mapping tool names to tool classes
            (e.g., agent.llm_tools_map)

    Returns:
        List of DoneSequence objects
    """
    return [parse_done_sequence(seq, tools_map) for seq in sequences]
</file>

<file path="langroid/agent/xml_tool_message.py">
import re
from collections.abc import Mapping
from typing import Any, Dict, List, Optional, Union, get_args, get_origin

from lxml import etree
from pydantic import BaseModel, ConfigDict

from langroid.agent.tool_message import ToolMessage

# For Union type handling - check if we have Python 3.10+ UnionType
HAS_UNION_TYPE = False
try:
    from types import UnionType  # noqa: F401 # Used conditionally

    HAS_UNION_TYPE = True
except ImportError:
    pass


class XMLToolMessage(ToolMessage):
    """
    Abstract class for tools formatted using XML instead of JSON.

    When a subclass defines a field with the attribute `verbatim=True`,
    instructions are sent to the LLM to ensure the field's content is:
        - preserved as is, including whitespace, indents, quotes, newlines, etc
            with no escaping, and
        - enclosed in a CDATA section in the XML output.
    This is useful for LLMs sending code as part of a tool;
    results can be far superior compared to sending code in JSON-formatted tools,
    where code needs to confirm to JSON's strict rules and escaping requirements.
    (see test_xml_tool_message.py for an example).

    """

    request: str
    purpose: str

    _allow_llm_use: bool = True

    model_config = ConfigDict(
        # Inherit settings from ToolMessage
        extra="allow",
        arbitrary_types_allowed=False,
        validate_default=True,
        validate_assignment=True,
        json_schema_extra={"exclude": ["purpose", "id"]},
    )

    # XMLToolMessage-specific settings as class methods to avoid Pydantic
    # treating them as model fields
    @classmethod
    def _get_excluded_fields(cls) -> set[str]:
        return {"purpose", "id"}

    # Root element for XML formatting
    @classmethod
    def _get_root_element(cls) -> str:
        return "tool"

    @classmethod
    def extract_field_values(cls, formatted_string: str) -> Optional[Dict[str, Any]]:
        """
        Extracts field values from an XML-formatted string.

        Args:
            formatted_string (str): The XML-formatted string to parse.

        Returns:
            Optional[Dict[str, Any]]: A dictionary containing the extracted field
                values, where keys are the XML element names and values are their
                corresponding contents.
            Returns None if parsing fails or the root element is not a dictionary.

        Raises:
            etree.XMLSyntaxError: If the input string is not valid XML.
        """
        # SECURITY: Initialize XMLParser with flags to prevent
        # XML External Entity (XXE), billion laughs, and external DTD attacks by
        # disabling entity resolution, DTD loading, and network access;
        # `strip_cdata=False` is needed to preserve
        # content within CDATA sections (e.g., for code).
        parser = etree.XMLParser(
            strip_cdata=False,
            resolve_entities=False,
            load_dtd=False,
            no_network=True,
        )
        root = etree.fromstring(formatted_string.encode("utf-8"), parser=parser)

        def parse_element(element: etree._Element) -> Any:
            # Skip elements starting with underscore
            if element.tag.startswith("_"):
                return {}

            field_info = cls.model_fields.get(element.tag)
            is_verbatim = (
                field_info
                and hasattr(field_info, "json_schema_extra")
                and field_info.json_schema_extra is not None
                and isinstance(field_info.json_schema_extra, dict)
                and field_info.json_schema_extra.get("verbatim", False)
            )

            if is_verbatim:
                # For code elements, preserve the content as is, including whitespace
                content = element.text if element.text else ""
                # Strip leading and trailing triple backticks if present,
                # accounting for whitespace
                return (
                    content.strip().removeprefix("```").removesuffix("```").strip()
                    if content.strip().startswith("```")
                    and content.strip().endswith("```")
                    else content
                )
            elif len(element) == 0:
                # For non-code leaf elements, strip whitespace
                return element.text.strip() if element.text else ""
            else:
                # For branch elements, handle potential lists or nested structures
                children = [parse_element(child) for child in element]
                if all(child.tag == element[0].tag for child in element):
                    # If all children have the same tag, treat as a list
                    return children
                else:
                    # Otherwise, treat as a dictionary
                    result = {child.tag: parse_element(child) for child in element}
                    # Check if this corresponds to a nested Pydantic model
                    if (
                        field_info
                        and isinstance(field_info.annotation, type)
                        and issubclass(field_info.annotation, BaseModel)
                    ):
                        return field_info.annotation(**result)
                    return result

        result = parse_element(root)
        if not isinstance(result, dict):
            return None
        # Filter out empty dictionaries from skipped underscore fields
        return {k: v for k, v in result.items() if v != {}}

    @classmethod
    def parse(cls, formatted_string: str) -> Optional["XMLToolMessage"]:
        """
        Parses the XML-formatted string and returns an instance of the class.

        Args:
            formatted_string (str): The XML-formatted string to parse.

        Returns:
            Optional["XMLToolMessage"]: An instance of the class if parsing succeeds,
                None otherwise.
        """
        try:
            parsed_data = cls.extract_field_values(formatted_string)
            if parsed_data is None:
                return None

            # Use Pydantic's parse_obj to create and validate the instance
            return cls.model_validate(parsed_data)
        except Exception as e:
            from langroid.exceptions import XMLException

            raise XMLException(f"Error parsing XML: {str(e)}")

    @classmethod
    def find_verbatim_fields(
        cls, prefix: str = "", parent_cls: Optional[type[BaseModel]] = None
    ) -> List[str]:
        verbatim_fields = []
        for field_name, field_info in (parent_cls or cls).model_fields.items():
            full_name = f"{prefix}.{field_name}" if prefix else field_name
            if (
                hasattr(field_info, "json_schema_extra")
                and field_info.json_schema_extra is not None
                and isinstance(field_info.json_schema_extra, dict)
                and field_info.json_schema_extra.get("verbatim", False)
            ) or field_name == "code":
                verbatim_fields.append(full_name)
            if isinstance(field_info.annotation, type) and issubclass(
                field_info.annotation, BaseModel
            ):
                verbatim_fields.extend(
                    cls.find_verbatim_fields(full_name, field_info.annotation)
                )
        return verbatim_fields

    @classmethod
    def format_instructions(cls, tool: bool = False) -> str:
        fields = [
            f for f in cls.model_fields.keys() if f not in cls._get_excluded_fields()
        ]

        instructions = """
        To use this tool, please provide the required information in an XML-like 
        format. Here's how to structure your input:\n\n
        """

        preamble = "Placeholders:\n"
        xml_format = f"Formatting example:\n\n<{cls._get_root_element()}>\n"

        def format_field(
            field_name: str,
            field_type: Any,
            indent: str = "",
            path: str = "",
        ) -> None:
            nonlocal preamble, xml_format
            current_path = f"{path}.{field_name}" if path else field_name

            origin = get_origin(field_type)
            args = get_args(field_type)

            # Handle Union types (including Optional types like List[Person] | None)
            # Support both typing.Union and types.UnionType (Python 3.10+ | syntax)
            is_union = origin is Union
            if HAS_UNION_TYPE:
                from types import UnionType as _UnionType

                is_union = is_union or origin is _UnionType

            if is_union:
                # Filter out None type for Optional types
                non_none_args = [arg for arg in args if arg is not type(None)]
                if len(non_none_args) == 1:
                    # This is an Optional type, process the non-None type
                    field_type = non_none_args[0]
                    origin = get_origin(field_type)
                    args = get_args(field_type)
                # If there are multiple non-None types, fall through to default handling

            if (
                origin is None
                and isinstance(field_type, type)
                and issubclass(field_type, BaseModel)
            ):
                preamble += (
                    f"{field_name.upper()} = [nested structure for {field_name}]\n"
                )
                xml_format += f"{indent}<{field_name}>\n"
                for sub_field, sub_field_info in field_type.model_fields.items():
                    format_field(
                        sub_field,
                        sub_field_info.annotation,
                        indent + "  ",
                        current_path,
                    )
                xml_format += f"{indent}</{field_name}>\n"
            elif origin in (list, List) or (field_type is list):
                item_type = args[0] if args else Any
                if isinstance(item_type, type) and issubclass(item_type, BaseModel):
                    preamble += (
                        f"{field_name.upper()} = "
                        f"[list of nested structures for {field_name}]\n"
                    )
                else:
                    preamble += (
                        f"{field_name.upper()} = "
                        f"[list of {getattr(item_type, '__name__', str(item_type))} "
                        f"for {field_name}]\n"
                    )
                xml_format += f"{indent}<{field_name}>\n"
                xml_format += (
                    f"{indent}  <item>"
                    f"[{getattr(item_type, '__name__', str(item_type))} value]"
                    f"</item>\n"
                )
                xml_format += f"{indent}  ...\n"
                xml_format += f"{indent}</{field_name}>\n"
            elif origin in (dict, Dict) or (
                isinstance(field_type, type) and issubclass(field_type, Mapping)
            ):
                key_type, value_type = args if len(args) == 2 else (Any, Any)
                preamble += (
                    f"{field_name.upper()} = "
                    f"[dictionary with "
                    f"{getattr(key_type, '__name__', str(key_type))} keys and "
                    f"{getattr(value_type, '__name__', str(value_type))} values]\n"
                )
                xml_format += f"{indent}<{field_name}>\n"
                xml_format += (
                    f"{indent}  <{getattr(key_type, '__name__', str(key_type))}>"
                    f"[{getattr(value_type, '__name__', str(value_type))} value]"
                    f"</{getattr(key_type, '__name__', str(key_type))}>\n"
                )
                xml_format += f"{indent}  ...\n"
                xml_format += f"{indent}</{field_name}>\n"
            else:
                preamble += f"{field_name.upper()} = [value for {field_name}]\n"
                if current_path in verbatim_fields:
                    xml_format += (
                        f"{indent}<{field_name}>"
                        f"<![CDATA[{{{field_name.upper()}}}]]></{field_name}>\n"
                    )
                else:
                    xml_format += (
                        f"{indent}<{field_name}>"
                        f"{{{field_name.upper()}}}</{field_name}>\n"
                    )

        verbatim_fields = cls.find_verbatim_fields()

        for field in fields:
            field_info = cls.model_fields[field]
            field_type = field_info.annotation
            # Ensure we have a valid type
            if field_type is None:
                continue
            format_field(field, field_type)

        xml_format += f"</{cls._get_root_element()}>"

        verbatim_alert = ""
        if len(verbatim_fields) > 0:
            verbatim_alert = f"""
            EXTREMELY IMPORTANT: For these fields:
            {', '.join(verbatim_fields)},
            the contents MUST be wrapped in a CDATA section, and the content
            must be written verbatim WITHOUT any modifications or escaping,
            such as spaces, tabs, indents, newlines, quotes, etc.
            """

        examples_str = ""
        if cls.examples():
            examples_str = "EXAMPLES:\n" + cls.usage_examples()

        return f"""
            TOOL: {cls.default_value("request")}
            PURPOSE: {cls.default_value("purpose")} 

            {instructions}
            {preamble}
            {xml_format}

            Make sure to replace the placeholders with actual values 
            when using the tool.                
            {verbatim_alert}            
            {examples_str}
            """.lstrip()

    def format_example(self) -> str:
        """
        Format the current instance as an XML example.

        Returns:
            str: A string representation of the current instance in XML format.

        Raises:
            ValueError: If the result from etree.tostring is not a string.
        """

        def create_element(
            parent: etree._Element, name: str, value: Any, path: str = ""
        ) -> None:
            if value is None:
                return

            elem = etree.SubElement(parent, name)
            current_path = f"{path}.{name}" if path else name

            if isinstance(value, list):
                for item in value:
                    create_element(elem, "item", item, current_path)
            elif isinstance(value, dict):
                for k, v in value.items():
                    create_element(elem, k, v, current_path)
            elif isinstance(value, BaseModel):
                # Handle nested Pydantic models
                for field_name, field_value in value.model_dump().items():
                    create_element(elem, field_name, field_value, current_path)
            else:
                if current_path in self.__class__.find_verbatim_fields():
                    elem.text = etree.CDATA(str(value))
                else:
                    elem.text = str(value)

        root = etree.Element(self._get_root_element())
        exclude_fields: set[str] = self._get_excluded_fields()
        for name, value in self.model_dump().items():
            if name not in exclude_fields:
                create_element(root, name, value)

        result = etree.tostring(root, encoding="unicode", pretty_print=True)
        if not isinstance(result, str):
            raise ValueError("Unexpected non-string result from etree.tostring")
        return result

    @classmethod
    def find_candidates(cls, text: str) -> List[str]:
        """
        Finds XML-like tool message candidates in text, with relaxed opening tag rules.

        Args:
            text: Input text to search for XML structures.

        Returns:
            List of XML strings. For fragments missing the root opening tag but having
            valid XML structure and root closing tag, prepends the root opening tag.

        Example:
            With root_tag="tool", given:
            "Hello <field1>data</field1> </tool>"
            Returns: ["<tool><field1>data</field1></tool>"]
        """

        root_tag = cls._get_root_element()
        opening_tag = f"<{root_tag}>"
        closing_tag = f"</{root_tag}>"

        candidates = []
        pos = 0
        while True:
            # Look for either proper opening tag or closing tag
            start_normal = text.find(opening_tag, pos)
            end = text.find(closing_tag, pos)

            if start_normal == -1 and end == -1:
                break

            if start_normal != -1:
                # Handle normal case (has opening tag)
                end = text.find(closing_tag, start_normal)
                if end != -1:
                    candidates.append(text[start_normal : end + len(closing_tag)])
                    pos = max(end + len(closing_tag), start_normal + 1)
                    continue
                elif start_normal == text.rfind(opening_tag):
                    # last fragment - ok to miss closing tag
                    candidates.append(text[start_normal:] + closing_tag)
                    return candidates
                else:
                    pos = start_normal + 1
                    continue

            if end != -1:
                # Look backwards for first XML tag
                text_before = text[pos:end]
                first_tag_match = re.search(r"<\w+>", text_before)
                if first_tag_match:
                    start = pos + first_tag_match.start()
                    candidates.append(
                        opening_tag + text[start : end + len(closing_tag)]
                    )
                pos = end + len(closing_tag)

        return candidates
</file>

<file path="langroid/cachedb/__init__.py">
from . import base

from . import redis_cachedb

__all__ = [
    "base",
    "redis_cachedb",
]
</file>

<file path="langroid/cachedb/base.py">
from abc import ABC, abstractmethod
from typing import Any, Dict, List

from pydantic_settings import BaseSettings


class CacheDBConfig(BaseSettings):
    """Configuration model for CacheDB."""

    pass


class CacheDB(ABC):
    """Abstract base class for a cache database."""

    @abstractmethod
    def store(self, key: str, value: Any) -> None:
        """
        Abstract method to store a value associated with a key.

        Args:
            key (str): The key under which to store the value.
            value (Any): The value to store.
        """
        pass

    @abstractmethod
    def retrieve(self, key: str) -> Dict[str, Any] | str | None:
        """
        Abstract method to retrieve the value associated with a key.

        Args:
            key (str): The key to retrieve the value for.

        Returns:
            dict: The value associated with the key.
        """
        pass

    @abstractmethod
    def delete_keys(self, keys: List[str]) -> None:
        """
        Delete the keys from the cache.

        Args:
            keys (List[str]): The keys to delete.
        """
        pass

    @abstractmethod
    def delete_keys_pattern(self, pattern: str) -> None:
        """
        Delete all keys with the given pattern

        Args:
            prefix (str): The pattern to match.
        """
        pass
</file>

<file path="langroid/cachedb/redis_cachedb.py">
import json
import logging
import os
from contextlib import AbstractContextManager, contextmanager
from typing import Any, Dict, List, TypeVar

import fakeredis
import redis
from dotenv import load_dotenv

from langroid.cachedb.base import CacheDB, CacheDBConfig

T = TypeVar("T", bound="RedisCache")
logger = logging.getLogger(__name__)


class RedisCacheConfig(CacheDBConfig):
    """Configuration model for RedisCache."""

    fake: bool = False


class RedisCache(CacheDB):
    """Redis implementation of the CacheDB."""

    _warned_password: bool = False

    def __init__(self, config: RedisCacheConfig):
        """
        Initialize a RedisCache with the given config.

        Args:
            config (RedisCacheConfig): The configuration to use.
        """
        self.config = config
        load_dotenv()

        if self.config.fake:
            self.pool = fakeredis.FakeStrictRedis()  # type: ignore
        else:
            redis_password = os.getenv("REDIS_PASSWORD")
            redis_host = os.getenv("REDIS_HOST") or None
            redis_port = os.getenv("REDIS_PORT")
            if None in [redis_password, redis_host, redis_port]:
                if not RedisCache._warned_password:
                    logger.warning(
                        """REDIS_PASSWORD, REDIS_HOST, REDIS_PORT not set in .env file,
                        using fake redis client"""
                    )
                    RedisCache._warned_password = True
                self.pool = fakeredis.FakeStrictRedis()  # type: ignore
            else:
                self.pool = redis.ConnectionPool(  # type: ignore
                    host=redis_host,
                    port=redis_port,
                    password=redis_password,
                    max_connections=500,
                    socket_timeout=5,
                    socket_keepalive=True,
                    retry_on_timeout=True,
                    health_check_interval=30,
                )

    @contextmanager  # type: ignore
    def redis_client(self) -> AbstractContextManager[T]:  # type: ignore
        """Cleanly open and close a redis client, avoids max clients exceeded error"""
        if isinstance(self.pool, fakeredis.FakeStrictRedis):
            yield self.pool
        else:
            client: T = redis.Redis(connection_pool=self.pool)
            try:
                yield client
            finally:
                client.close()

    def close_all_connections(self) -> None:
        with self.redis_client() as client:  # type: ignore
            clients = client.client_list()
            for c in clients:
                client.client_kill(c["addr"])

    def clear(self) -> None:
        """Clear keys from current db."""
        with self.redis_client() as client:  # type: ignore
            client.flushdb()

    def clear_all(self) -> None:
        """Clear all keys from all dbs."""
        with self.redis_client() as client:  # type: ignore
            client.flushall()

    def store(self, key: str, value: Any) -> None:
        """
        Store a value associated with a key.

        Args:
            key (str): The key under which to store the value.
            value (Any): The value to store.
        """
        with self.redis_client() as client:  # type: ignore
            try:
                client.set(key, json.dumps(value))
            except redis.exceptions.ConnectionError:
                logger.warning("Redis connection error, not storing key/value")
                return None

    def retrieve(self, key: str) -> Dict[str, Any] | str | None:
        """
        Retrieve the value associated with a key.

        Args:
            key (str): The key to retrieve the value for.

        Returns:
            dict|str|None: The value associated with the key.
        """
        with self.redis_client() as client:  # type: ignore
            try:
                value = client.get(key)
            except redis.exceptions.ConnectionError:
                logger.warning("Redis connection error, returning None")
                return None
            return json.loads(value) if value else None

    def delete_keys(self, keys: List[str]) -> None:
        """
        Delete the keys from the cache.

        Args:
            keys (List[str]): The keys to delete.
        """
        with self.redis_client() as client:  # type: ignore
            try:
                client.delete(*keys)
            except redis.exceptions.ConnectionError:
                logger.warning("Redis connection error, not deleting keys")
                return None

    def delete_keys_pattern(self, pattern: str) -> None:
        """
        Delete the keys matching the pattern from the cache.

        Args:
            prefix (str): The pattern to match.
        """
        with self.redis_client() as client:  # type: ignore
            try:
                keys = client.keys(pattern)
                if len(keys) > 0:
                    client.delete(*keys)
            except redis.exceptions.ConnectionError:
                logger.warning("Redis connection error, not deleting keys")
                return None
</file>

<file path="langroid/embedding_models/protoc/__init__.py">

</file>

<file path="langroid/embedding_models/protoc/embeddings_pb2_grpc.py">
# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
"""Client and server classes corresponding to protobuf-defined services."""
import grpc

import langroid.embedding_models.protoc.embeddings_pb2 as embeddings__pb2


class EmbeddingStub(object):
    """Missing associated documentation comment in .proto file."""

    def __init__(self, channel):
        """Constructor.

        Args:
            channel: A grpc.Channel.
        """
        self.Embed = channel.unary_unary(
            "/Embedding/Embed",
            request_serializer=embeddings__pb2.EmbeddingRequest.SerializeToString,
            response_deserializer=embeddings__pb2.BatchEmbeds.FromString,
        )


class EmbeddingServicer(object):
    """Missing associated documentation comment in .proto file."""

    def Embed(self, request, context):
        """Missing associated documentation comment in .proto file."""
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details("Method not implemented!")
        raise NotImplementedError("Method not implemented!")


def add_EmbeddingServicer_to_server(servicer, server):
    rpc_method_handlers = {
        "Embed": grpc.unary_unary_rpc_method_handler(
            servicer.Embed,
            request_deserializer=embeddings__pb2.EmbeddingRequest.FromString,
            response_serializer=embeddings__pb2.BatchEmbeds.SerializeToString,
        ),
    }
    generic_handler = grpc.method_handlers_generic_handler(
        "Embedding", rpc_method_handlers
    )
    server.add_generic_rpc_handlers((generic_handler,))


# This class is part of an EXPERIMENTAL API.
class Embedding(object):
    """Missing associated documentation comment in .proto file."""

    @staticmethod
    def Embed(
        request,
        target,
        options=(),
        channel_credentials=None,
        call_credentials=None,
        insecure=False,
        compression=None,
        wait_for_ready=None,
        timeout=None,
        metadata=None,
    ):
        return grpc.experimental.unary_unary(
            request,
            target,
            "/Embedding/Embed",
            embeddings__pb2.EmbeddingRequest.SerializeToString,
            embeddings__pb2.BatchEmbeds.FromString,
            options,
            channel_credentials,
            insecure,
            call_credentials,
            compression,
            wait_for_ready,
            timeout,
            metadata,
        )
</file>

<file path="langroid/embedding_models/protoc/embeddings_pb2.py">
# -*- coding: utf-8 -*-
# Generated by the protocol buffer compiler.  DO NOT EDIT!
# source: embeddings.proto
# Protobuf Python Version: 4.25.1
"""Generated protocol buffer code."""
from google.protobuf import descriptor as _descriptor
from google.protobuf import descriptor_pool as _descriptor_pool
from google.protobuf import symbol_database as _symbol_database
from google.protobuf.internal import builder as _builder

# @@protoc_insertion_point(imports)

_sym_db = _symbol_database.Default()


DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
    b'\n\x10\x65mbeddings.proto"K\n\x10\x45mbeddingRequest\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x12\n\nbatch_size\x18\x02 \x01(\x05\x12\x0f\n\x07strings\x18\x03 \x03(\t"%\n\x0b\x42\x61tchEmbeds\x12\x16\n\x06\x65mbeds\x18\x01 \x03(\x0b\x32\x06.Embed"\x16\n\x05\x45mbed\x12\r\n\x05\x65mbed\x18\x01 \x03(\x02\x32\x37\n\tEmbedding\x12*\n\x05\x45mbed\x12\x11.EmbeddingRequest\x1a\x0c.BatchEmbeds"\x00\x62\x06proto3'
)

_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "embeddings_pb2", _globals)
if _descriptor._USE_C_DESCRIPTORS == False:
    DESCRIPTOR._options = None
    _globals["_EMBEDDINGREQUEST"]._serialized_start = 20
    _globals["_EMBEDDINGREQUEST"]._serialized_end = 95
    _globals["_BATCHEMBEDS"]._serialized_start = 97
    _globals["_BATCHEMBEDS"]._serialized_end = 134
    _globals["_EMBED"]._serialized_start = 136
    _globals["_EMBED"]._serialized_end = 158
    _globals["_EMBEDDING"]._serialized_start = 160
    _globals["_EMBEDDING"]._serialized_end = 215
# @@protoc_insertion_point(module_scope)
</file>

<file path="langroid/embedding_models/protoc/embeddings_pb2.pyi">
from typing import (
    ClassVar as _ClassVar,
)
from typing import (
    Iterable as _Iterable,
)
from typing import (
    Mapping as _Mapping,
)
from typing import (
    Optional as _Optional,
)
from typing import (
    Union as _Union,
)

from google.protobuf import descriptor as _descriptor
from google.protobuf import message as _message
from google.protobuf.internal import containers as _containers

DESCRIPTOR: _descriptor.FileDescriptor

class EmbeddingRequest(_message.Message):
    __slots__ = ("model_name", "batch_size", "strings")
    MODEL_NAME_FIELD_NUMBER: _ClassVar[int]
    BATCH_SIZE_FIELD_NUMBER: _ClassVar[int]
    STRINGS_FIELD_NUMBER: _ClassVar[int]
    model_name: str
    batch_size: int
    strings: _containers.RepeatedScalarFieldContainer[str]
    def __init__(
        self,
        model_name: _Optional[str] = ...,
        batch_size: _Optional[int] = ...,
        strings: _Optional[_Iterable[str]] = ...,
    ) -> None: ...

class BatchEmbeds(_message.Message):
    __slots__ = ("embeds",)
    EMBEDS_FIELD_NUMBER: _ClassVar[int]
    embeds: _containers.RepeatedCompositeFieldContainer[Embed]
    def __init__(
        self, embeds: _Optional[_Iterable[_Union[Embed, _Mapping]]] = ...
    ) -> None: ...

class Embed(_message.Message):
    __slots__ = ("embed",)
    EMBED_FIELD_NUMBER: _ClassVar[int]
    embed: _containers.RepeatedScalarFieldContainer[float]
    def __init__(self, embed: _Optional[_Iterable[float]] = ...) -> None: ...
</file>

<file path="langroid/embedding_models/protoc/embeddings.proto">
syntax = "proto3";

service Embedding {
    rpc Embed (EmbeddingRequest) returns (BatchEmbeds) {};
}

message EmbeddingRequest {
    string model_name = 1;
    int32 batch_size = 2;
    repeated string strings = 3;
}

message BatchEmbeds {
    repeated Embed embeds = 1;
}

message Embed {
    repeated float embed = 1;
}
</file>

<file path="langroid/embedding_models/__init__.py">
from . import base
from . import models
from . import remote_embeds

from .base import (
    EmbeddingModel,
    EmbeddingModelsConfig,
)
from .models import (
    OpenAIEmbeddings,
    OpenAIEmbeddingsConfig,
    SentenceTransformerEmbeddings,
    SentenceTransformerEmbeddingsConfig,
    LlamaCppServerEmbeddings,
    LlamaCppServerEmbeddingsConfig,
    GeminiEmbeddings,
    GeminiEmbeddingsConfig,
    embedding_model,
)
from .remote_embeds import (
    RemoteEmbeddingsConfig,
    RemoteEmbeddings,
)


__all__ = [
    "base",
    "models",
    "remote_embeds",
    "EmbeddingModel",
    "EmbeddingModelsConfig",
    "OpenAIEmbeddings",
    "OpenAIEmbeddingsConfig",
    "SentenceTransformerEmbeddings",
    "SentenceTransformerEmbeddingsConfig",
    "LlamaCppServerEmbeddings",
    "LlamaCppServerEmbeddingsConfig",
    "GeminiEmbeddings",
    "GeminiEmbeddingsConfig",
    "embedding_model",
    "RemoteEmbeddingsConfig",
    "RemoteEmbeddings",
]
</file>

<file path="langroid/embedding_models/base.py">
import logging
from abc import ABC, abstractmethod

import numpy as np
from pydantic_settings import BaseSettings

from langroid.mytypes import EmbeddingFunction

logging.getLogger("openai").setLevel(logging.ERROR)


class EmbeddingModelsConfig(BaseSettings):
    model_type: str = "openai"
    dims: int = 0
    context_length: int = 512
    batch_size: int = 512


class EmbeddingModel(ABC):
    """
    Abstract base class for an embedding model.
    """

    def clone(self) -> "EmbeddingModel":
        """
        Return a copy of this embedding model suitable for use in cloned agents.
        Default behaviour attempts to deep-copy the model configuration and
        instantiate a fresh model of the same type; if that is not possible,
        the original instance is reused.
        """
        config = getattr(self, "config", None)
        if config is not None and hasattr(config, "model_copy"):
            try:
                return type(self)(config.model_copy(deep=True))  # type: ignore[call-arg]
            except Exception:
                pass
        return self

    @classmethod
    def create(cls, config: EmbeddingModelsConfig) -> "EmbeddingModel":
        from langroid.embedding_models.models import (
            AzureOpenAIEmbeddings,
            AzureOpenAIEmbeddingsConfig,
            FastEmbedEmbeddings,
            FastEmbedEmbeddingsConfig,
            GeminiEmbeddings,
            GeminiEmbeddingsConfig,
            LlamaCppServerEmbeddings,
            LlamaCppServerEmbeddingsConfig,
            OpenAIEmbeddings,
            OpenAIEmbeddingsConfig,
            SentenceTransformerEmbeddings,
            SentenceTransformerEmbeddingsConfig,
        )
        from langroid.embedding_models.remote_embeds import (
            RemoteEmbeddings,
            RemoteEmbeddingsConfig,
        )

        if isinstance(config, RemoteEmbeddingsConfig):
            return RemoteEmbeddings(config)
        elif isinstance(config, OpenAIEmbeddingsConfig):
            return OpenAIEmbeddings(config)
        elif isinstance(config, AzureOpenAIEmbeddingsConfig):
            return AzureOpenAIEmbeddings(config)
        elif isinstance(config, SentenceTransformerEmbeddingsConfig):
            return SentenceTransformerEmbeddings(config)
        elif isinstance(config, FastEmbedEmbeddingsConfig):
            return FastEmbedEmbeddings(config)
        elif isinstance(config, LlamaCppServerEmbeddingsConfig):
            return LlamaCppServerEmbeddings(config)
        elif isinstance(config, GeminiEmbeddingsConfig):
            return GeminiEmbeddings(config)
        else:
            raise ValueError(f"Unknown embedding config: {config.__class__.__name__}")

    @abstractmethod
    def embedding_fn(self) -> EmbeddingFunction:
        pass

    @property
    @abstractmethod
    def embedding_dims(self) -> int:
        pass

    def similarity(self, text1: str, text2: str) -> float:
        """Compute cosine similarity between two texts."""
        [emb1, emb2] = self.embedding_fn()([text1, text2])
        return float(
            np.array(emb1)
            @ np.array(emb2)
            / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
        )
</file>

<file path="langroid/embedding_models/models.py">
import atexit
import os
from functools import cached_property
from typing import Any, Callable, Dict, List, Optional

import requests
import tiktoken
from dotenv import load_dotenv
from openai import AzureOpenAI, OpenAI
from pydantic_settings import SettingsConfigDict

from langroid.embedding_models.base import EmbeddingModel, EmbeddingModelsConfig
from langroid.exceptions import LangroidImportError
from langroid.language_models.provider_params import LangDBParams
from langroid.mytypes import Embeddings
from langroid.parsing.utils import batched

AzureADTokenProvider = Callable[[], str]


class OpenAIEmbeddingsConfig(EmbeddingModelsConfig):
    model_type: str = "openai"
    model_name: str = "text-embedding-3-small"
    api_key: str = ""
    api_base: Optional[str] = None
    organization: str = ""
    dims: int = 1536
    context_length: int = 8192
    langdb_params: LangDBParams = LangDBParams()

    model_config = SettingsConfigDict(env_prefix="OPENAI_")


class AzureOpenAIEmbeddingsConfig(EmbeddingModelsConfig):
    model_type: str = "azure-openai"
    model_name: str = "text-embedding-3-small"
    api_key: str = ""
    api_base: str = ""
    deployment_name: Optional[str] = None
    # api_version defaulted to 2024-06-01 as per https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/embeddings?tabs=python-new
    # change this to required  supported version
    api_version: Optional[str] = "2024-06-01"
    # TODO: Add auth support for Azure OpenAI via AzureADTokenProvider
    azure_ad_token: Optional[str] = None
    azure_ad_token_provider: Optional[AzureADTokenProvider] = None
    dims: int = 1536
    context_length: int = 8192

    model_config = SettingsConfigDict(env_prefix="AZURE_OPENAI_")


class SentenceTransformerEmbeddingsConfig(EmbeddingModelsConfig):
    model_type: str = "sentence-transformer"
    model_name: str = "BAAI/bge-large-en-v1.5"
    context_length: int = 512
    data_parallel: bool = False
    # Select device (e.g. "cuda", "cpu") when data parallel is disabled
    device: Optional[str] = None
    # Select devices when data parallel is enabled
    devices: Optional[list[str]] = None


class FastEmbedEmbeddingsConfig(EmbeddingModelsConfig):
    """Config for qdrant/fastembed embeddings,
    see here: https://github.com/qdrant/fastembed
    """

    model_type: str = "fastembed"
    model_name: str = "BAAI/bge-small-en-v1.5"
    batch_size: int = 256
    cache_dir: Optional[str] = None
    threads: Optional[int] = None
    parallel: Optional[int] = None
    additional_kwargs: Dict[str, Any] = {}


class LlamaCppServerEmbeddingsConfig(EmbeddingModelsConfig):
    api_base: str = ""
    context_length: int = 2048
    batch_size: int = 2048


class GeminiEmbeddingsConfig(EmbeddingModelsConfig):
    model_type: str = "gemini"
    model_name: str = "models/text-embedding-004"
    api_key: str = ""
    dims: int = 768
    batch_size: int = 512


class EmbeddingFunctionCallable:
    """
    A callable class designed to generate embeddings for a list of texts using
    the OpenAI or Azure OpenAI API, with automatic retries on failure.

    Attributes:
        embed_model (EmbeddingModel): An instance of EmbeddingModel that provides
               configuration and utilities for generating embeddings.

    Methods:
        __call__(input: List[str]) -> Embeddings: Generate embeddings for
                                a list of input texts.
    """

    def __init__(self, embed_model: EmbeddingModel, batch_size: int = 512):
        """
        Initialize the EmbeddingFunctionCallable with a specific model.

        Args:
            model ( OpenAIEmbeddings or AzureOpenAIEmbeddings): An instance of
                            OpenAIEmbeddings or AzureOpenAIEmbeddings to use for
                            generating embeddings.
            batch_size (int): Batch size
        """
        self.embed_model = embed_model
        self.batch_size = batch_size

    def __call__(self, input: List[str]) -> Embeddings:
        """
        Generate embeddings for a given list of input texts using the OpenAI API,
        with retries on failure.

        This method:
        - Truncates each text in the input list to the model's maximum context length.
        - Processes the texts in batches to generate embeddings efficiently.
        - Automatically retries the embedding generation process with exponential
        backoff in case of failures.

        Args:
            input (List[str]): A list of input texts to generate embeddings for.

        Returns:
            Embeddings: A list of embedding vectors corresponding to the input texts.
        """
        embeds = []
        if isinstance(self.embed_model, (OpenAIEmbeddings, AzureOpenAIEmbeddings)):
            # Truncate texts to context length while preserving text format
            truncated_texts = self.embed_model.truncate_texts(input)

            # Process in batches
            for batch in batched(truncated_texts, self.batch_size):
                result = self.embed_model.client.embeddings.create(
                    input=batch, model=self.embed_model.config.model_name  # type: ignore
                )
                batch_embeds = [d.embedding for d in result.data]
                embeds.extend(batch_embeds)

        elif isinstance(self.embed_model, SentenceTransformerEmbeddings):
            if self.embed_model.config.data_parallel:
                embeds = self.embed_model.model.encode_multi_process(
                    input,
                    self.embed_model.pool,
                    batch_size=self.batch_size,
                ).tolist()
            else:
                for str_batch in batched(input, self.batch_size):
                    batch_embeds = self.embed_model.model.encode(
                        str_batch, convert_to_numpy=True
                    ).tolist()  # type: ignore
                    embeds.extend(batch_embeds)

        elif isinstance(self.embed_model, FastEmbedEmbeddings):
            embeddings = self.embed_model.model.embed(
                input, batch_size=self.batch_size, parallel=self.embed_model.parallel
            )

            embeds = [embedding.tolist() for embedding in embeddings]
        elif isinstance(self.embed_model, LlamaCppServerEmbeddings):
            for input_string in input:
                tokenized_text = self.embed_model.tokenize_string(input_string)
                for token_batch in batched(tokenized_text, self.batch_size):
                    gen_embedding = self.embed_model.generate_embedding(
                        self.embed_model.detokenize_string(list(token_batch))
                    )
                    embeds.append(gen_embedding)
        elif isinstance(self.embed_model, GeminiEmbeddings):
            embeds = self.embed_model.generate_embeddings(input)
        return embeds


class OpenAIEmbeddings(EmbeddingModel):
    def __init__(self, config: OpenAIEmbeddingsConfig = OpenAIEmbeddingsConfig()):
        super().__init__()
        self.config = config
        load_dotenv()

        # Check if using LangDB
        self.is_langdb = self.config.model_name.startswith("langdb/")

        if self.is_langdb:
            self.config.model_name = self.config.model_name.replace("langdb/", "")
            self.config.api_base = self.config.langdb_params.base_url
            project_id = self.config.langdb_params.project_id
            if project_id:
                self.config.api_base += "/" + project_id + "/v1"
            self.config.api_key = self.config.langdb_params.api_key

        if not self.config.api_key:
            self.config.api_key = os.getenv("OPENAI_API_KEY", "")

        self.config.organization = os.getenv("OPENAI_ORGANIZATION", "")

        if self.config.api_key == "":
            if self.is_langdb:
                raise ValueError(
                    """
                    LANGDB_API_KEY must be set in .env or your environment 
                    to use OpenAIEmbeddings via LangDB.
                    """
                )
            else:
                raise ValueError(
                    """
                    OPENAI_API_KEY must be set in .env or your environment 
                    to use OpenAIEmbeddings.
                    """
                )

        self.client = OpenAI(
            base_url=self.config.api_base,
            api_key=self.config.api_key,
            organization=self.config.organization,
        )
        model_for_tokenizer = self.config.model_name
        if model_for_tokenizer.startswith("openai/"):
            self.config.model_name = model_for_tokenizer.replace("openai/", "")
        self.tokenizer = tiktoken.encoding_for_model(self.config.model_name)

    def truncate_texts(self, texts: List[str]) -> List[str] | List[List[int]]:
        """
        Truncate texts to the embedding model's context length.
        TODO: Maybe we should show warning, and consider doing T5 summarization?
        """
        truncated_tokens = [
            self.tokenizer.encode(text, disallowed_special=())[
                : self.config.context_length
            ]
            for text in texts
        ]

        if self.is_langdb:
            # LangDB embedding endpt only works with strings, not tokens
            return [self.tokenizer.decode(tokens) for tokens in truncated_tokens]
        return truncated_tokens

    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
        return EmbeddingFunctionCallable(self, self.config.batch_size)

    @property
    def embedding_dims(self) -> int:
        return self.config.dims


class AzureOpenAIEmbeddings(EmbeddingModel):
    """
    Azure OpenAI embeddings model implementation.
    """

    def __init__(
        self, config: AzureOpenAIEmbeddingsConfig = AzureOpenAIEmbeddingsConfig()
    ):
        """
        Initializes Azure OpenAI embeddings model.

        Args:
            config: Configuration for Azure OpenAI embeddings model.
        Raises:
            ValueError: If required Azure config values are not set.
        """
        super().__init__()
        self.config = config
        load_dotenv()

        if self.config.api_key == "":
            raise ValueError(
                """AZURE_OPENAI_API_KEY env variable must be set to use 
            AzureOpenAIEmbeddings. Please set the AZURE_OPENAI_API_KEY value 
            in your .env file."""
            )

        if self.config.api_base == "":
            raise ValueError(
                """AZURE_OPENAI_API_BASE env variable must be set to use 
            AzureOpenAIEmbeddings. Please set the AZURE_OPENAI_API_BASE value 
            in your .env file."""
            )
        self.client = AzureOpenAI(
            api_key=self.config.api_key,
            api_version=self.config.api_version,
            azure_endpoint=self.config.api_base,
            azure_deployment=self.config.deployment_name,
        )
        self.tokenizer = tiktoken.encoding_for_model(self.config.model_name)

    def truncate_texts(self, texts: List[str]) -> List[str] | List[List[int]]:
        """
        Truncate texts to the embedding model's context length.
        TODO: Maybe we should show warning, and consider doing T5 summarization?
        """
        return [
            self.tokenizer.encode(text, disallowed_special=())[
                : self.config.context_length
            ]
            for text in texts
        ]

    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
        """Get the embedding function for Azure OpenAI.

        Returns:
            Callable that generates embeddings for input texts.
        """
        return EmbeddingFunctionCallable(self, self.config.batch_size)

    @property
    def embedding_dims(self) -> int:
        return self.config.dims


STEC = SentenceTransformerEmbeddingsConfig


class SentenceTransformerEmbeddings(EmbeddingModel):
    def __init__(self, config: STEC = STEC()):
        # this is an "extra" optional dependency, so we import it here
        try:
            from sentence_transformers import SentenceTransformer
            from transformers import AutoTokenizer
        except ImportError:
            raise ImportError(
                """
                To use sentence_transformers embeddings, 
                you must install langroid with the [hf-embeddings] extra, e.g.:
                pip install "langroid[hf-embeddings]"
                """
            )

        super().__init__()
        self.config = config

        self.model = SentenceTransformer(
            self.config.model_name,
            device=self.config.device,
        )
        if self.config.data_parallel:
            self.pool = self.model.start_multi_process_pool(
                self.config.devices  # type: ignore
            )
            atexit.register(
                lambda: SentenceTransformer.stop_multi_process_pool(self.pool)
            )

        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
        self.config.context_length = self.tokenizer.model_max_length

    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
        return EmbeddingFunctionCallable(self, self.config.batch_size)

    @property
    def embedding_dims(self) -> int:
        dims = self.model.get_sentence_embedding_dimension()
        if dims is None:
            raise ValueError(
                f"Could not get embedding dimension for model {self.config.model_name}"
            )
        return dims  # type: ignore


class FastEmbedEmbeddings(EmbeddingModel):
    def __init__(self, config: FastEmbedEmbeddingsConfig = FastEmbedEmbeddingsConfig()):
        try:
            from fastembed import TextEmbedding
        except ImportError:
            raise LangroidImportError("fastembed", extra="fastembed")

        super().__init__()
        self.config = config
        self.batch_size = config.batch_size
        self.parallel = config.parallel

        self.model = TextEmbedding(
            model_name=self.config.model_name,
            cache_dir=self.config.cache_dir,
            threads=self.config.threads,
            **self.config.additional_kwargs,
        )

    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
        return EmbeddingFunctionCallable(self, self.config.batch_size)

    @cached_property
    def embedding_dims(self) -> int:
        embed_func = self.embedding_fn()
        return len(embed_func(["text"])[0])


LCSEC = LlamaCppServerEmbeddingsConfig


class LlamaCppServerEmbeddings(EmbeddingModel):
    def __init__(self, config: LCSEC = LCSEC()):
        super().__init__()
        self.config = config

        if self.config.api_base == "":
            raise ValueError(
                """Api Base MUST be set for Llama Server Embeddings.
                """
            )

        self.tokenize_url = self.config.api_base + "/tokenize"
        self.detokenize_url = self.config.api_base + "/detokenize"
        self.embedding_url = self.config.api_base + "/embeddings"

    def tokenize_string(self, text: str) -> List[int]:
        data = {"content": text, "add_special": False, "with_pieces": False}
        response = requests.post(self.tokenize_url, json=data)

        if response.status_code == 200:
            tokens = response.json()["tokens"]
            if not (isinstance(tokens, list) and isinstance(tokens[0], (int, float))):
                # not all(isinstance(token, (int, float)) for token in tokens):
                raise ValueError(
                    """Tokenizer endpoint has not returned the correct format. 
                   Is the URL correct?
                """
                )
            return tokens
        else:
            raise requests.HTTPError(
                self.tokenize_url,
                response.status_code,
                "Failed to connect to tokenization provider",
            )

    def detokenize_string(self, tokens: List[int]) -> str:
        data = {"tokens": tokens}
        response = requests.post(self.detokenize_url, json=data)

        if response.status_code == 200:
            text = response.json()["content"]
            if not isinstance(text, str):
                raise ValueError(
                    """Deokenizer endpoint has not returned the correct format. 
                   Is the URL correct?
                """
                )
            return text
        else:
            raise requests.HTTPError(
                self.detokenize_url,
                response.status_code,
                "Failed to connect to detokenization provider",
            )

    def truncate_string_to_context_size(self, text: str) -> str:
        tokens = self.tokenize_string(text)
        tokens = tokens[: self.config.context_length]
        return self.detokenize_string(tokens)

    def generate_embedding(self, text: str) -> List[int | float]:
        data = {"content": text}
        response = requests.post(self.embedding_url, json=data)

        if response.status_code == 200:
            embeddings = self._extract_embedding(response.json())
            if not (
                isinstance(embeddings, list) and isinstance(embeddings[0], (int, float))
            ):
                raise ValueError(
                    """Embedding endpoint has not returned the correct format.
                   Is the URL correct?
                """
                )
            return embeddings
        else:
            raise requests.HTTPError(
                self.embedding_url,
                response.status_code,
                "Failed to connect to embedding provider",
            )

    def _extract_embedding(
        self, response_json: dict[str, Any] | list[Any]
    ) -> List[int | float]:
        """
        Extract embedding vector from llama.cpp response.

        Handles multiple response formats:
        1. Native /embedding: {"embedding": [floats]}
        2. Array format: [{"embedding": [floats]}]
        3. Double-nested: [{"embedding": [[floats]]}]
        4. OpenAI /v1/embeddings: {"data": [{"embedding": [floats]}]}
        5. Nested in dict: {"embedding": [[floats]]}

        Args:
            response_json: The JSON response from llama.cpp server

        Returns:
            List of floats representing the embedding vector

        Raises:
            ValueError: If response format is not recognized
        """
        import json

        # Try native format first: {"embedding": [floats]}
        if isinstance(response_json, dict) and "embedding" in response_json:
            embeddings = response_json["embedding"]
            # Check if it's [floats]
            if isinstance(embeddings, list) and len(embeddings) > 0:
                if isinstance(embeddings[0], (int, float)):
                    return embeddings
                # Might be nested: {"embedding": [[floats]]}
                if isinstance(embeddings[0], list) and len(embeddings[0]) > 0:
                    if isinstance(embeddings[0][0], (int, float)):
                        return embeddings[0]

        # Try OpenAI format: {"data": [{"embedding": [floats]}]}
        if isinstance(response_json, dict) and "data" in response_json:
            data = response_json["data"]
            if isinstance(data, list) and len(data) > 0:
                if isinstance(data[0], dict) and "embedding" in data[0]:
                    embeddings = data[0]["embedding"]
                    if isinstance(embeddings, list) and len(embeddings) > 0:
                        if isinstance(embeddings[0], (int, float)):
                            return embeddings

        # Try array format: [{"embedding": [floats]}] or [{"embedding": [[floats]]}]
        if isinstance(response_json, list) and len(response_json) > 0:
            first_item = response_json[0]
            if isinstance(first_item, dict) and "embedding" in first_item:
                embeddings = first_item["embedding"]
                # Check if it's [floats]
                if isinstance(embeddings, list) and len(embeddings) > 0:
                    if isinstance(embeddings[0], (int, float)):
                        return embeddings
                    # Check if it's [[floats]]
                    if isinstance(embeddings[0], list) and len(embeddings[0]) > 0:
                        if isinstance(embeddings[0][0], (int, float)):
                            return embeddings[0]

        raise ValueError(
            f"Unsupported embedding response format from {self.embedding_url}. "
            f"Response: {json.dumps(response_json)[:500]}"
        )

    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
        return EmbeddingFunctionCallable(self, self.config.batch_size)

    @property
    def embedding_dims(self) -> int:
        return self.config.dims


class GeminiEmbeddings(EmbeddingModel):
    def __init__(self, config: GeminiEmbeddingsConfig = GeminiEmbeddingsConfig()):
        try:
            from google import genai
        except ImportError as e:
            raise LangroidImportError(extra="google-genai", error=str(e))
        super().__init__()
        self.config = config
        load_dotenv()
        self.config.api_key = os.getenv("GEMINI_API_KEY", "")

        if self.config.api_key == "":
            raise ValueError(
                """
                GEMINI_API_KEY env variable must be set to use GeminiEmbeddings.
                """
            )
        self.client = genai.Client(api_key=self.config.api_key)

    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
        return EmbeddingFunctionCallable(self, self.config.batch_size)

    def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Generates embeddings for a list of input texts."""
        all_embeddings: List[List[float]] = []

        for batch in batched(texts, self.config.batch_size):
            result = self.client.models.embed_content(  # type: ignore[attr-defined]
                model=self.config.model_name,
                contents=batch,  # type: ignore
            )

            if not hasattr(result, "embeddings") or not isinstance(
                result.embeddings, list
            ):
                raise ValueError(
                    "Unexpected format for embeddings: missing or incorrect type"
                )

            # Extract .values from ContentEmbedding objects
            all_embeddings.extend(
                [emb.values for emb in result.embeddings]  # type: ignore
            )

        return all_embeddings

    @property
    def embedding_dims(self) -> int:
        return self.config.dims


def embedding_model(embedding_fn_type: str = "openai") -> EmbeddingModel:
    """
    Args:
        embedding_fn_type: Type of embedding model to use. Options are:
         - "openai",
         - "azure-openai",
         - "sentencetransformer", or
         - "fastembed".
            (others may be added in the future)
    Returns:
        EmbeddingModel: The corresponding embedding model class.
    """
    if embedding_fn_type == "openai":
        return OpenAIEmbeddings  # type: ignore
    elif embedding_fn_type == "azure-openai":
        return AzureOpenAIEmbeddings  # type: ignore
    elif embedding_fn_type == "fastembed":
        return FastEmbedEmbeddings  # type: ignore
    elif embedding_fn_type == "llamacppserver":
        return LlamaCppServerEmbeddings  # type: ignore
    elif embedding_fn_type == "gemini":
        return GeminiEmbeddings  # type: ignore
    else:  # default sentence transformer
        return SentenceTransformerEmbeddings  # type: ignore
</file>

<file path="langroid/embedding_models/remote_embeds.py">
"""
If run as a script, starts an RPC server which handles remote
embedding requests:

For example:
python3 -m langroid.embedding_models.remote_embeds --port `port`

where `port` is the port at which the service is exposed.  Currently,
supports insecure connections only, and this should NOT be exposed to
the internet.
"""

import atexit
import subprocess
import time
from typing import Callable, Optional

import grpc
from fire import Fire

import langroid.embedding_models.models as em
import langroid.embedding_models.protoc.embeddings_pb2 as embeddings_pb
import langroid.embedding_models.protoc.embeddings_pb2_grpc as embeddings_grpc
from langroid.mytypes import Embeddings


class RemoteEmbeddingRPCs(embeddings_grpc.EmbeddingServicer):
    def __init__(
        self,
        model_name: str,
        batch_size: int,
        data_parallel: bool,
        device: Optional[str],
        devices: Optional[list[str]],
    ):
        super().__init__()
        self.embedding_fn = em.SentenceTransformerEmbeddings(
            em.SentenceTransformerEmbeddingsConfig(
                model_name=model_name,
                batch_size=batch_size,
                data_parallel=data_parallel,
                device=device,
                devices=devices,
            )
        ).embedding_fn()

    def Embed(
        self, request: embeddings_pb.EmbeddingRequest, _: grpc.RpcContext
    ) -> embeddings_pb.BatchEmbeds:
        embeds = self.embedding_fn(list(request.strings))

        embeds_pb = [embeddings_pb.Embed(embed=e) for e in embeds]

        return embeddings_pb.BatchEmbeds(embeds=embeds_pb)


class RemoteEmbeddingsConfig(em.SentenceTransformerEmbeddingsConfig):
    api_base: str = "localhost"
    port: int = 50052
    # The below are used only when waiting for server creation
    poll_delay: float = 0.01
    max_retries: int = 1000


class RemoteEmbeddings(em.SentenceTransformerEmbeddings):
    def __init__(self, config: RemoteEmbeddingsConfig = RemoteEmbeddingsConfig()):
        super().__init__(config)
        self.config: RemoteEmbeddingsConfig = config
        self.have_started_server: bool = False

    def embedding_fn(self) -> Callable[[list[str]], Embeddings]:
        def fn(texts: list[str]) -> Embeddings:
            url = f"{self.config.api_base}:{self.config.port}"
            with grpc.insecure_channel(url) as channel:
                stub = embeddings_grpc.EmbeddingStub(channel)  # type: ignore
                response = stub.Embed(
                    embeddings_pb.EmbeddingRequest(
                        strings=texts,
                    )
                )

                return [list(emb.embed) for emb in response.embeds]

        def with_handling(texts: list[str]) -> Embeddings:
            # In local mode, start the server if it has not already
            # been started
            if self.config.api_base == "localhost" and not self.have_started_server:
                try:
                    return fn(texts)
                # Occurs when the server hasn't been started
                except grpc.RpcError:
                    self.have_started_server = True
                    # Start the server
                    proc = subprocess.Popen(
                        [
                            "python3",
                            __file__,
                            "--bind_address_base",
                            self.config.api_base,
                            "--port",
                            str(self.config.port),
                            "--batch_size",
                            str(self.config.batch_size),
                            "--model_name",
                            self.config.model_name,
                        ],
                    )

                    atexit.register(lambda: proc.terminate())

                    for _ in range(self.config.max_retries - 1):
                        try:
                            return fn(texts)
                        except grpc.RpcError:
                            time.sleep(self.config.poll_delay)

            # The remote is not local or we have exhausted retries
            # We should now raise an error if the server is not accessible
            return fn(texts)

        return with_handling


async def serve(
    bind_address_base: str = "localhost",
    port: int = 50052,
    batch_size: int = 512,
    data_parallel: bool = False,
    device: Optional[str] = None,
    devices: Optional[list[str]] = None,
    model_name: str = "BAAI/bge-large-en-v1.5",
) -> None:
    """Starts the RPC server."""
    server = grpc.aio.server()
    embeddings_grpc.add_EmbeddingServicer_to_server(
        RemoteEmbeddingRPCs(
            model_name=model_name,
            batch_size=batch_size,
            data_parallel=data_parallel,
            device=device,
            devices=devices,
        ),
        server,
    )  # type: ignore
    url = f"{bind_address_base}:{port}"
    server.add_insecure_port(url)
    await server.start()
    print(f"Embedding server started, listening on {url}")
    await server.wait_for_termination()


if __name__ == "__main__":
    Fire(serve)
</file>

<file path="langroid/language_models/prompt_formatter/__init__.py">
from . import base
from . import llama2_formatter
from .base import PromptFormatter
from .llama2_formatter import Llama2Formatter
from ..config import PromptFormatterConfig
from ..config import Llama2FormatterConfig


__all__ = [
    "PromptFormatter",
    "Llama2Formatter",
    "PromptFormatterConfig",
    "Llama2FormatterConfig",
    "base",
    "llama2_formatter",
]
</file>

<file path="langroid/language_models/prompt_formatter/base.py">
import logging
from abc import ABC, abstractmethod
from typing import List

from langroid.language_models.base import LLMMessage
from langroid.language_models.config import PromptFormatterConfig

logger = logging.getLogger(__name__)


class PromptFormatter(ABC):
    """
    Abstract base class for a prompt formatter
    """

    def __init__(self, config: PromptFormatterConfig):
        self.config = config

    @staticmethod
    def create(formatter: str) -> "PromptFormatter":
        from langroid.language_models.config import HFPromptFormatterConfig
        from langroid.language_models.prompt_formatter.hf_formatter import HFFormatter

        return HFFormatter(HFPromptFormatterConfig(model_name=formatter))

    @abstractmethod
    def format(self, messages: List[LLMMessage]) -> str:
        """
        Convert sequence of messages (system, user, assistant, user, assistant...user)
            to a single prompt formatted according to the specific format type,
            to be used in a /completions endpoint.

        Args:
            messages (List[LLMMessage]): chat history as a sequence of messages

        Returns:
            (str): formatted version of chat history

        """
        pass
</file>

<file path="langroid/language_models/prompt_formatter/hf_formatter.py">
"""
Prompt formatter based on HuggingFace `AutoTokenizer.apply_chat_template` method
from their Transformers library. It searches the hub for a model matching the
specified name, and uses the first one it finds. We assume that all matching
models will have the same tokenizer, so we just use the first one.
"""

import logging
import re
from typing import Any, List, Set, Tuple, Type

from jinja2.exceptions import TemplateError

from langroid.language_models.base import LanguageModel, LLMMessage, Role
from langroid.language_models.config import HFPromptFormatterConfig
from langroid.language_models.prompt_formatter.base import PromptFormatter

logger = logging.getLogger(__name__)


def try_import_hf_modules() -> Tuple[Type[Any], Type[Any]]:
    """
    Attempts to import the AutoTokenizer class from the transformers package.
    Returns:
        The AutoTokenizer class if successful.
    Raises:
        ImportError: If the transformers package is not installed.
    """
    try:
        from huggingface_hub import HfApi
        from transformers import AutoTokenizer

        return AutoTokenizer, HfApi
    except ImportError:
        raise ImportError(
            """
            You are trying to use some/all of:
            HuggingFace transformers.AutoTokenizer,
            huggingface_hub.HfApi,
            but these are not not installed 
            by default with Langroid. Please install langroid using the 
            `transformers` extra, like so:
            pip install "langroid[transformers]"
            or equivalent.
            """
        )


def find_hf_formatter(model_name: str) -> str:
    AutoTokenizer, HfApi = try_import_hf_modules()
    hf_api = HfApi()
    # try to find a matching model, with progressivly shorter prefixes of model_name
    model_name = model_name.lower().split("/")[-1]
    parts = re.split("[:\\-_]", model_name)
    parts = [p.lower() for p in parts if p != ""]
    for i in range(len(parts), 0, -1):
        prefix = "-".join(parts[:i])
        models = hf_api.list_models(
            task="text-generation",
            model_name=prefix,
        )
        try:
            mdl = next(models)
            tokenizer = AutoTokenizer.from_pretrained(mdl.id)
            if tokenizer.chat_template is not None:
                return str(mdl.id)
            else:
                continue
        except Exception:
            continue

    return ""


class HFFormatter(PromptFormatter):
    models: Set[str] = set()  # which models have been used for formatting

    def __init__(self, config: HFPromptFormatterConfig):
        super().__init__(config)
        AutoTokenizer, HfApi = try_import_hf_modules()
        self.config: HFPromptFormatterConfig = config
        hf_api = HfApi()
        models = hf_api.list_models(
            task="text-generation",
            model_name=config.model_name,
        )
        try:
            mdl = next(models)
        except StopIteration:
            raise ValueError(f"Model {config.model_name} not found on HuggingFace Hub")

        self.tokenizer = AutoTokenizer.from_pretrained(mdl.id)
        if self.tokenizer.chat_template is None:
            raise ValueError(
                f"Model {config.model_name} does not support chat template"
            )
        elif mdl.id not in HFFormatter.models:
            # only warn if this is the first time we've used this mdl.id
            logger.warning(
                f"""
            Using HuggingFace {mdl.id} for prompt formatting: 
            This is the CHAT TEMPLATE. If this is not what you intended,
            consider specifying a more complete model name for the formatter.
             
            {self.tokenizer.chat_template}
            """
            )
        HFFormatter.models.add(mdl.id)

    def format(self, messages: List[LLMMessage]) -> str:
        sys_msg, chat_msgs, user_msg = LanguageModel.get_chat_history_components(
            messages
        )
        # build msg dicts expected by AutoTokenizer.apply_chat_template
        sys_msg_dict = dict(role=Role.SYSTEM.value, content=sys_msg)
        chat_dicts = []
        for user, assistant in chat_msgs:
            chat_dicts.append(dict(role=Role.USER.value, content=user))
            chat_dicts.append(dict(role=Role.ASSISTANT.value, content=assistant))
        chat_dicts.append(dict(role=Role.USER.value, content=user_msg))
        all_dicts = [sys_msg_dict] + chat_dicts
        try:
            # apply chat template
            result = self.tokenizer.apply_chat_template(all_dicts, tokenize=False)
        except TemplateError:
            # this likely means the model doesn't support a system msg,
            # so combine it with the first user msg
            first_user_msg = chat_msgs[0][0] if len(chat_msgs) > 0 else user_msg
            first_user_msg = sys_msg + "\n\n" + first_user_msg
            chat_dicts[0] = dict(role=Role.USER.value, content=first_user_msg)
            result = self.tokenizer.apply_chat_template(chat_dicts, tokenize=False)
        return str(result)
</file>

<file path="langroid/language_models/prompt_formatter/llama2_formatter.py">
import logging
from typing import List, Tuple

from langroid.language_models.base import LanguageModel, LLMMessage
from langroid.language_models.config import Llama2FormatterConfig
from langroid.language_models.prompt_formatter.base import PromptFormatter

logger = logging.getLogger(__name__)


BOS: str = "<s>"
EOS: str = "</s>"
B_INST: str = "[INST]"
E_INST: str = "[/INST]"
B_SYS: str = "<<SYS>>\n"
E_SYS: str = "\n<</SYS>>\n\n"
SPECIAL_TAGS: List[str] = [B_INST, E_INST, BOS, EOS, "<<SYS>>", "<</SYS>>"]


class Llama2Formatter(PromptFormatter):
    def __int__(self, config: Llama2FormatterConfig) -> None:
        super().__init__(config)
        self.config: Llama2FormatterConfig = config

    def format(self, messages: List[LLMMessage]) -> str:
        sys_msg, chat_msgs, user_msg = LanguageModel.get_chat_history_components(
            messages
        )
        return self._get_prompt_from_components(sys_msg, chat_msgs, user_msg)

    def _get_prompt_from_components(
        self,
        system_prompt: str,
        chat_history: List[Tuple[str, str]],
        user_message: str,
    ) -> str:
        """
        For llama2 models, convert chat history into a single
        prompt for Llama2 models, for use in the /completions endpoint
        (as opposed to the /chat/completions endpoint).
        See:
        https://www.reddit.com/r/LocalLLaMA/comments/155po2p/get_llama_2_prompt_format_right/
        https://github.com/facebookresearch/llama/blob/main/llama/generation.py#L44

        Args:
            system_prompt (str): system prompt, typically specifying role/task.
            chat_history (List[Tuple[str,str]]): List of (user, assistant) pairs
            user_message (str): user message, at the end of the chat, i.e. the message
                for which we want to generate a response.

        Returns:
            str: Prompt for Llama2 models

        Typical structure of the formatted prompt:
        Note important that the first [INST], [/INST] surrounds the system prompt,
        together with the first user message. A lot of libs seem to miss this detail.

        <s>[INST] <<SYS>>
        You are are a helpful... bla bla.. assistant
        <</SYS>>

        Hi there! [/INST] Hello! How can I help you today? </s><s>[INST]
        What is a neutron star? [/INST] A neutron star is a ... </s><s>
        [INST] Okay cool, thank you! [/INST] You're welcome! </s><s>
        [INST] Ah, I have one more question.. [/INST]
        """
        bos = BOS if self.config.use_bos_eos else ""
        eos = EOS if self.config.use_bos_eos else ""
        text = f"{bos}{B_INST} {B_SYS}{system_prompt}{E_SYS}"
        for user_input, response in chat_history:
            text += (
                f"{user_input.strip()} {E_INST} {response.strip()} {eos}{bos} {B_INST} "
            )
        text += f"{user_message.strip()} {E_INST}"
        return text
</file>

<file path="langroid/language_models/azure_openai.py">
import logging
from typing import Callable

from dotenv import load_dotenv
from httpx import Timeout
from openai import AsyncAzureOpenAI, AzureOpenAI
from pydantic_settings import SettingsConfigDict

from langroid.language_models.openai_gpt import (
    OpenAIGPT,
    OpenAIGPTConfig,
)

azureStructuredOutputList = [
    "2024-08-06",
    "2024-11-20",
]

azureStructuredOutputAPIMin = "2024-08-01-preview"

logger = logging.getLogger(__name__)


class AzureConfig(OpenAIGPTConfig):
    """
    Configuration for Azure OpenAI GPT.

    Attributes:
        type (str): should be ``azure.``
        api_version (str): can be set in the ``.env`` file as
            ``AZURE_OPENAI_API_VERSION.``
        deployment_name (str|None): can be optionally set in the ``.env`` file as
            ``AZURE_OPENAI_DEPLOYMENT_NAME`` and should be based the custom name you
            chose for your deployment when you deployed a model.
        model_name (str): [DEPRECATED] can be set in the ``.env``
            file as ``AZURE_OPENAI_MODEL_NAME``
            and should be based on the model name chosen during setup.
        chat_model (str): the chat model name to use. Can be set via
            the env variable ``AZURE_OPENAI_CHAT_MODEL``.
            Recommended to use this instead of ``model_name``.

    """

    api_key: str = ""  # CAUTION: set this ONLY via env var AZURE_OPENAI_API_KEY
    type: str = "azure"
    api_version: str = "2023-05-15"
    deployment_name: str | None = None
    model_name: str = ""
    api_base: str = ""

    # Alternatively, bring your own clients:
    azure_openai_client_provider: Callable[[], AzureOpenAI] | None = None
    azure_openai_async_client_provider: Callable[[], AsyncAzureOpenAI] | None = None

    # all of the vars above can be set via env vars,
    # by upper-casing the name and prefixing with `env_prefix`, e.g.
    # AZURE_OPENAI_API_VERSION=2023-05-15
    # This is either done in the .env file, or via an explicit
    # `export AZURE_OPENAI_API_VERSION=...`
    model_config = SettingsConfigDict(env_prefix="AZURE_OPENAI_")

    def __init__(self, **kwargs) -> None:  # type: ignore
        if "model_name" in kwargs and "chat_model" not in kwargs:
            kwargs["chat_model"] = kwargs["model_name"]
        super().__init__(**kwargs)


class AzureGPT(OpenAIGPT):
    """
    Class to access OpenAI LLMs via Azure. These env variables can be obtained from the
    file `.azure_env`. Azure OpenAI doesn't support ``completion``
    """

    def __init__(self, config: AzureConfig):
        # This will auto-populate config values from .env file
        load_dotenv()
        super().__init__(config)
        self.config: AzureConfig = config

        if (
            self.config.azure_openai_client_provider
            or self.config.azure_openai_async_client_provider
        ):
            if not self.config.azure_openai_client_provider:
                self.client = None
                logger.warning(
                    "Using user-provided Azure OpenAI client, but only async "
                    "client has been provided. Synchronous calls will fail."
                )
            if not self.config.azure_openai_async_client_provider:
                self.async_client = None
                logger.warning(
                    "Using user-provided Azure OpenAI client, but no async "
                    "client has been provided. Asynchronous calls will fail."
                )

            if self.config.azure_openai_client_provider:
                self.client = self.config.azure_openai_client_provider()
            if self.config.azure_openai_async_client_provider:
                self.async_client = self.config.azure_openai_async_client_provider()
                self.async_client.timeout = Timeout(self.config.timeout)
        else:
            if self.config.api_key == "":
                raise ValueError(
                    """
                    AZURE_OPENAI_API_KEY not set in .env file,
                    please set it to your Azure API key."""
                )

            if self.config.api_base == "":
                raise ValueError(
                    """
                    AZURE_OPENAI_API_BASE not set in .env file,
                    please set it to your Azure API key."""
                )

            self.client = AzureOpenAI(
                api_key=self.config.api_key,
                azure_endpoint=self.config.api_base,
                api_version=self.config.api_version,
                azure_deployment=self.config.deployment_name,
            )
            self.async_client = AsyncAzureOpenAI(
                api_key=self.config.api_key,
                azure_endpoint=self.config.api_base,
                api_version=self.config.api_version,
                azure_deployment=self.config.deployment_name,
                timeout=Timeout(self.config.timeout),
            )

        self.supports_json_schema = (
            self.config.api_version >= azureStructuredOutputAPIMin
            and self.config.api_version in azureStructuredOutputList
        )
</file>

<file path="langroid/language_models/config.py">
from pydantic_settings import BaseSettings, SettingsConfigDict


class PromptFormatterConfig(BaseSettings):
    type: str = "llama2"

    model_config = SettingsConfigDict(env_prefix="FORMAT_", case_sensitive=False)


class Llama2FormatterConfig(PromptFormatterConfig):
    use_bos_eos: bool = False


class HFPromptFormatterConfig(PromptFormatterConfig):
    type: str = "hf"
    model_name: str
</file>

<file path="langroid/language_models/provider_params.py">
"""
Provider-specific parameter configurations for various LLM providers.
"""

from typing import Any, Dict, Optional

from pydantic_settings import BaseSettings, SettingsConfigDict

# Constants
LANGDB_BASE_URL = "https://api.us-east-1.langdb.ai"
PORTKEY_BASE_URL = "https://api.portkey.ai"
DUMMY_API_KEY = "xxx"


class LangDBParams(BaseSettings):
    """
    Parameters specific to LangDB integration.
    """

    api_key: str = DUMMY_API_KEY
    project_id: str = ""
    label: Optional[str] = None
    run_id: Optional[str] = None
    thread_id: Optional[str] = None
    base_url: str = LANGDB_BASE_URL

    model_config = SettingsConfigDict(env_prefix="LANGDB_")


class PortkeyParams(BaseSettings):
    """
    Parameters specific to Portkey integration.

    Portkey is an AI gateway that provides a unified API for multiple LLM providers,
    with features like automatic retries, fallbacks, load balancing, and observability.

    Example usage:
        # Use Portkey with Anthropic
        config = OpenAIGPTConfig(
            chat_model="portkey/anthropic/claude-3-sonnet-20240229",
            portkey_params=PortkeyParams(
                api_key="your-portkey-api-key",
                provider="anthropic"
            )
        )
    """

    api_key: str = DUMMY_API_KEY  # Portkey API key
    provider: str = ""  # Required: e.g., "openai", "anthropic", "cohere", etc.
    virtual_key: Optional[str] = None  # Optional: virtual key for the provider
    trace_id: Optional[str] = None  # Optional: trace ID for request tracking
    metadata: Optional[Dict[str, Any]] = None  # Optional: metadata for logging
    retry: Optional[Dict[str, Any]] = None  # Optional: retry configuration
    cache: Optional[Dict[str, Any]] = None  # Optional: cache configuration
    cache_force_refresh: Optional[bool] = None  # Optional: force cache refresh
    user: Optional[str] = None  # Optional: user identifier
    organization: Optional[str] = None  # Optional: organization identifier
    custom_headers: Optional[Dict[str, str]] = None  # Optional: additional headers
    base_url: str = PORTKEY_BASE_URL

    model_config = SettingsConfigDict(env_prefix="PORTKEY_")

    def get_headers(self) -> Dict[str, str]:
        """Generate Portkey-specific headers from parameters."""
        import json
        import os

        headers = {}

        if self.api_key and self.api_key != DUMMY_API_KEY:
            headers["x-portkey-api-key"] = self.api_key
        else:
            portkey_key = os.getenv("PORTKEY_API_KEY", "")
            if portkey_key:
                headers["x-portkey-api-key"] = portkey_key

        if self.provider:
            headers["x-portkey-provider"] = self.provider

        if self.virtual_key:
            headers["x-portkey-virtual-key"] = self.virtual_key

        if self.trace_id:
            headers["x-portkey-trace-id"] = self.trace_id

        if self.metadata:
            headers["x-portkey-metadata"] = json.dumps(self.metadata)

        if self.retry:
            headers["x-portkey-retry"] = json.dumps(self.retry)

        if self.cache:
            headers["x-portkey-cache"] = json.dumps(self.cache)

        if self.cache_force_refresh is not None:
            headers["x-portkey-cache-force-refresh"] = str(
                self.cache_force_refresh
            ).lower()

        if self.user:
            headers["x-portkey-user"] = self.user

        if self.organization:
            headers["x-portkey-organization"] = self.organization

        if self.custom_headers:
            headers.update(self.custom_headers)

        return headers

    def parse_model_string(self, model_string: str) -> tuple[str, str]:
        """
        Parse a model string like "portkey/anthropic/claude-3-sonnet"
        and extract provider and model name.

        Returns:
            tuple: (provider, model_name)
        """
        parts = model_string.split("/", 2)
        if len(parts) >= 3 and parts[0] == "portkey":
            _, provider, model = parts
            return provider, model
        else:
            model = model_string.replace("portkey/", "")
            return "", model

    def get_provider_api_key(
        self, provider: str, default_key: str = DUMMY_API_KEY
    ) -> str:
        """
        Get the API key for the provider from environment variables.

        Args:
            provider: The provider name (e.g., "anthropic", "openai")
            default_key: Default key to return if not found

        Returns:
            The API key for the provider
        """
        import os

        env_patterns = [
            f"{provider.upper()}_API_KEY",
            f"{provider.upper()}_KEY",
        ]

        for pattern in env_patterns:
            key = os.getenv(pattern, "")
            if key:
                return key

        return default_key
</file>

<file path="langroid/parsing/__init__.py">
from . import parser
from . import agent_chats
from . import code_parser
from . import document_parser
from . import parse_json
from . import para_sentence_split
from . import repo_loader
from . import url_loader
from . import table_loader
from . import urls
from . import utils
from . import search
from . import web_search

from .parser import (
    Splitter,
    MarkitdownXLSParsingConfig,
    MarkitdownXLSXParsingConfig,
    MarkitdownPPTXParsingConfig,
    PdfParsingConfig,
    DocxParsingConfig,
    DocParsingConfig,
    ParsingConfig,
    Parser,
)

__all__ = [
    "parser",
    "agent_chats",
    "code_parser",
    "document_parser",
    "parse_json",
    "para_sentence_split",
    "repo_loader",
    "url_loader",
    "table_loader",
    "urls",
    "utils",
    "search",
    "web_search",
    "Splitter",
    "PdfParsingConfig",
    "DocxParsingConfig",
    "DocParsingConfig",
    "ParsingConfig",
    "MarkitdownXLSXParsingConfig",
    "MarkitdownXLSParsingConfig",
    "MarkitdownPPTXParsingConfig",
    "Parser",
]

try:
    from . import spider

    spider
    __all__.append("spider")
except ImportError:
    pass
</file>

<file path="langroid/parsing/agent_chats.py">
from typing import Tuple, no_type_check

from pyparsing import Empty, Literal, ParseException, SkipTo, StringEnd, Word, alphanums


@no_type_check
def parse_message(msg: str) -> Tuple[str, str]:
    """
    Parse the intended recipient and content of a message.
    Message format is assumed to be TO[<recipient>]:<message>.
    The TO[<recipient>]: part is optional.

    Args:
        msg (str): message to parse

    Returns:
        str, str: task-name of intended recipient, and content of message
            (if recipient is not specified, task-name is empty string)

    """
    if msg is None:
        return "", ""

    # Grammar definition
    name = Word(alphanums)
    to_start = Literal("TO[").suppress()
    to_end = Literal("]:").suppress()
    to_field = (to_start + name("name") + to_end) | Empty().suppress()
    message = SkipTo(StringEnd())("text")

    # Parser definition
    parser = to_field + message

    try:
        parsed = parser.parse_string(msg)
        return parsed.name, parsed.text
    except ParseException:
        return "", msg
</file>

<file path="langroid/parsing/code_parser.py">
from functools import reduce
from typing import Callable, List

import tiktoken
from pydantic_settings import BaseSettings
from pygments import lex
from pygments.lexers import get_lexer_by_name
from pygments.token import Token

from langroid.mytypes import Document


def chunk_code(
    code: str, language: str, max_tokens: int, len_fn: Callable[[str], int]
) -> List[str]:
    """
    Chunk code into smaller pieces, so that we don't exceed the maximum
    number of tokens allowed by the embedding model.
    Args:
        code: string of code
        language: str as a file extension, e.g. "py", "yml"
        max_tokens: max tokens per chunk
        len_fn: function to get the length of a string in token units
    Returns:

    """
    lexer = get_lexer_by_name(language)
    tokens = list(lex(code, lexer))

    chunks = []
    current_chunk = ""
    for token_type, token_value in tokens:
        if token_type in Token.Text.Whitespace:
            current_chunk += token_value
        else:
            token_tokens = len_fn(token_value)
            if len_fn(current_chunk) + token_tokens <= max_tokens:
                current_chunk += token_value
            else:
                chunks.append(current_chunk)
                current_chunk = token_value

    if current_chunk:
        chunks.append(current_chunk)

    return chunks


class CodeParsingConfig(BaseSettings):
    extensions: List[str] = [
        "py",
        "java",
        "c",
        "cpp",
        "h",
        "hpp",
        "yml",
        "yaml",
        "toml",
        "cfg",  # e.g. setup.cfg
        "ini",
        "json",
        "rst",
        "sh",
        "bash",
    ]
    chunk_size: int = 500  # tokens
    token_encoding_model: str = "text-embedding-3-small"
    n_similar_docs: int = 4


class CodeParser:
    def __init__(self, config: CodeParsingConfig):
        self.config = config
        self.tokenizer = tiktoken.encoding_for_model(config.token_encoding_model)

    def num_tokens(self, text: str) -> int:
        """
        How many tokens are in the text, according to the tokenizer.
        This needs to be accurate, otherwise we may exceed the maximum
        number of tokens allowed by the model.
        Args:
            text: string to tokenize
        Returns:
            number of tokens in the text
        """
        tokens = self.tokenizer.encode(text)
        return len(tokens)

    def split(self, docs: List[Document]) -> List[Document]:
        """
        Split the documents into chunks, according to the config.splitter.
        Only the documents with a language in the config.extensions are split.
        !!! note
            We assume the metadata in each document has at least a `language` field,
            which is used to determine how to chunk the code.
        Args:
            docs: list of documents to split
        Returns:
            list of documents, where each document is a chunk; the metadata of the
            original document is duplicated for each chunk, so that when we retrieve a
            chunk, we immediately know info about the original document.
        """
        chunked_docs = [
            [
                Document(content=chunk, metadata=d.metadata)
                for chunk in chunk_code(
                    d.content,
                    d.metadata.language,  # type: ignore
                    self.config.chunk_size,
                    self.num_tokens,
                )
                if chunk.strip() != ""
            ]
            for d in docs
            if d.metadata.language in self.config.extensions  # type: ignore
        ]
        if len(chunked_docs) == 0:
            return []
        # collapse the list of lists into a single list
        return reduce(lambda x, y: x + y, chunked_docs)
</file>

<file path="langroid/parsing/file_attachment.py">
import base64
import mimetypes
import uuid
from pathlib import Path
from typing import Any, BinaryIO, Dict, Optional, Union
from urllib.parse import urlparse

from pydantic import BaseModel


class FileAttachment(BaseModel):
    """Represents a file attachment to be sent to an LLM API."""

    content: bytes
    filename: Optional[str] = None
    mime_type: str = "application/octet-stream"
    url: str | None = None
    detail: str | None = None

    def __init__(self, **data: Any) -> None:
        """Initialize with sensible defaults for filename if not provided."""
        if "filename" not in data or data["filename"] is None:
            # Generate a more readable unique filename
            unique_id = str(uuid.uuid4())[:8]
            data["filename"] = f"attachment_{unique_id}.bin"
        super().__init__(**data)

    @classmethod
    def _from_path(
        cls,
        file_path: Union[str, Path],
        detail: Optional[str] = None,
    ) -> "FileAttachment":
        """Create a FileAttachment from a file path.

        Args:
            file_path: Path to the file to attach

        Returns:
            FileAttachment instance
        """
        path = Path(file_path)
        with open(path, "rb") as f:
            content = f.read()

        mime_type, _ = mimetypes.guess_type(path)
        if mime_type is None:
            mime_type = "application/octet-stream"

        return cls(
            content=content,
            filename=path.name,
            mime_type=mime_type,
            detail=detail,
        )

    @classmethod
    def _from_url(
        cls,
        url: str,
        content: Optional[bytes] = None,
        filename: Optional[str] = None,
        mime_type: Optional[str] = None,
        detail: Optional[str] = None,
    ) -> "FileAttachment":
        """Create a FileAttachment from a URL.

        Args:
            url: URL to the file
            content: Optional raw bytes content (if already fetched)
            filename: Optional name to use for the file
            mime_type: MIME type of the content, guessed from filename or url

        Returns:
            FileAttachment instance
        """
        if filename is None and url:
            # Extract filename from URL if possible

            parsed_url = urlparse(url)
            path = parsed_url.path
            filename = path.split("/")[-1] if path else None

        if mime_type is None and filename:
            mime_type, _ = mimetypes.guess_type(filename)

        return cls(
            content=content or b"",  # Empty bytes if no content provided
            filename=filename,
            mime_type=mime_type or "application/octet-stream",
            url=url,
            detail=detail,
        )

    @classmethod
    def from_path(
        cls,
        path: Union[str, Path],
        detail: str | None = None,
    ) -> "FileAttachment":
        """Create a FileAttachment from either a local file path or a URL.

        Args:
            path_or_url: Path to the file or URL to fetch

        Returns:
            FileAttachment instance
        """
        # Convert to string if Path object
        path_str = str(path)

        # Check if it's a URL
        if path_str.startswith(("http://", "https://", "ftp://")):
            return cls._from_url(url=path_str, detail=detail)
        else:
            # Assume it's a local file path
            return cls._from_path(path_str, detail=detail)

    @classmethod
    def from_bytes(
        cls,
        content: bytes,
        filename: Optional[str] = None,
        mime_type: Optional[str] = None,
    ) -> "FileAttachment":
        """Create a FileAttachment from bytes content.

        Args:
            content: Raw bytes content
            filename: Optional name to use for the file
            mime_type: MIME type of the content, guessed from filename if provided

        Returns:
            FileAttachment instance
        """
        if mime_type is None and filename is not None:
            mime_type, _ = mimetypes.guess_type(filename)

        return cls(
            content=content,
            filename=filename,
            mime_type=mime_type or "application/octet-stream",
        )

    @classmethod
    def from_io(
        cls,
        file_obj: BinaryIO,
        filename: Optional[str] = None,
        mime_type: Optional[str] = None,
    ) -> "FileAttachment":
        """Create a FileAttachment from a file-like object.

        Args:
            file_obj: File-like object with binary content
            filename: Optional name to use for the file
            mime_type: MIME type of the content, guessed from filename if provided

        Returns:
            FileAttachment instance
        """
        content = file_obj.read()
        return cls.from_bytes(content, filename, mime_type)

    @classmethod
    def from_text(
        cls,
        text: str,
        filename: Optional[str] = None,
        mime_type: str = "text/plain",
        encoding: str = "utf-8",
    ) -> "FileAttachment":
        """Create a FileAttachment from text content.

        Args:
            text: Text content to include
            filename: Optional name to use for the file
            mime_type: MIME type of the content
            encoding: Text encoding to use

        Returns:
            FileAttachment instance
        """
        content = text.encode(encoding)
        return cls(content=content, filename=filename, mime_type=mime_type)

    def to_base64(self) -> str:
        """Convert content to base64 encoding.

        Returns:
            Base64 encoded string
        """
        return base64.b64encode(self.content).decode("utf-8")

    def to_data_uri(self) -> str:
        """Convert content to a data URI.

        Returns:
            A data URI string containing the base64-encoded content with MIME type
        """
        base64_content = self.to_base64()
        return f"data:{self.mime_type};base64,{base64_content}"

    def to_dict(self, model: str) -> Dict[str, Any]:
        """
        Convert to a dictionary suitable for API requests.
        Tested only for PDF files.

        Returns:
            Dictionary with file data
        """
        if (
            self.mime_type
            and self.mime_type.startswith("image/")
            or "gemini" in model.lower()
        ):
            # for gemini models, we use `image_url` for both pdf-files and images

            image_url_dict = {}

            # If we have a URL and it's a full http/https URL, use it directly
            if self.url and (
                self.url.startswith("http://") or self.url.startswith("https://")
            ):
                image_url_dict["url"] = self.url
            # Otherwise use base64 data URI
            else:
                image_url_dict["url"] = self.to_data_uri()

            # Add detail parameter if specified
            if self.detail:
                image_url_dict["detail"] = self.detail

            return dict(
                type="image_url",
                image_url=image_url_dict,
            )
        else:
            # For non-image files
            return dict(
                type="file",
                file=dict(
                    filename=self.filename,
                    file_data=self.to_data_uri(),
                ),
            )
</file>

<file path="langroid/parsing/md_parser.py">
import re
from typing import Any, List

from pydantic import BaseModel, Field, field_validator

HEADER_CONTEXT_SEP = "\n...\n"


# Pydantic model definition for a node in the markdown hierarchy
class Node(BaseModel):
    content: str  # The text of the header or content block
    path: List[str]  # List of header texts from root to this node
    children: List["Node"] = Field(default_factory=list)
    # Nested children nodes

    def __repr__(self) -> str:
        # for debug printing
        return (
            f"Node(content={self.content!r}, path={self.path!r}, "
            f"children={len(self.children)})"
        )

    # Pydantic v1 requires forward references for self-referencing models
    # Forward references will be resolved with the update_forward_refs call below.


# Resolve forward references for Node (required for recursive models in Pydantic)
Node.model_rebuild()


def _cleanup_text(text: str) -> str:
    # 1) Convert alternative newline representations (any CRLF or CR) to a single '\n'
    text = text.replace("\r\n", "\n").replace("\r", "\n")

    # 2) Replace 3 or more consecutive newlines with exactly 2 newlines
    text = re.sub(r"\n{3,}", "\n\n", text)

    return text


HEADING_RE = re.compile(r"^(#{1,6})\s+(.*)$")


def parse_markdown_headings(md_text: str) -> List[Node]:
    """
    Parse `md_text` to extract a heading-based hierarchy, skipping lines
    that look like headings inside fenced code blocks. Each heading node
    will have a child node for the text that appears between this heading
    and the next heading.

    Returns a list of top-level Node objects.

    Example structure:
        Node(content='# Chapter 1', path=['# Chapter 1'], children=[
            Node(content='Intro paragraph...', path=['# Chapter 1'], children=[]),
            Node(content='## Section 1.1', path=['# Chapter 1', '## Section 1.1'],
                 children=[
                  Node(content='Some text in Section 1.1.', path=[...], children=[])
            ]),
            ...
        ])
    """
    # If doc is empty or only whitespace, return []
    if not md_text.strip():
        return []

    lines = md_text.splitlines(True)  # keep the newline characters

    # We'll scan line-by-line, track code-fence status, collect headings
    headings = []  # list of (level, heading_line, start_line_idx)
    in_code_fence = False
    fence_marker = None  # track which triple-backtick or ~~~ opened

    for i, line in enumerate(lines):
        # Check if we're toggling in/out of a fenced code block
        # Typically triple backtick or triple tilde: ``` or ~~~
        # We do a *loose* check: a line that starts with at least 3 backticks or tildes
        # ignoring trailing text. You can refine as needed.
        fence_match = re.match(r"^(```+|~~~+)", line.strip())
        if fence_match:
            # If we are not in a fence, we enter one;
            # If we are in a fence, we exit if the marker matches
            marker = fence_match.group(1)  # e.g. "```" or "~~~~"
            if not in_code_fence:
                in_code_fence = True
                fence_marker = marker[:3]  # store triple backtick or triple tilde
            else:
                # only close if the fence_marker matches
                # E.g. if we opened with ```, we close only on ```
                if fence_marker and marker.startswith(fence_marker):
                    in_code_fence = False
                    fence_marker = None

        if not in_code_fence:
            # Check if the line is a heading
            m = HEADING_RE.match(line)
            if m:
                hashes = m.group(1)  # e.g. "##"
                heading_text = line.rstrip("\n")  # entire line, exact
                level = len(hashes)
                headings.append((level, heading_text, i))

    # If no headings found, return a single root node with the entire text
    if not headings:
        return [Node(content=md_text.strip(), path=[], children=[])]

    # Add a sentinel heading at the end-of-file, so we can slice the last block
    # after the final real heading. We'll use level=0 so it doesn't form a real node.
    headings.append((0, "", len(lines)))

    # Now we build "heading blocks" with
    # (level, heading_text, start_line, end_line, content)
    heading_blocks = []
    for idx in range(len(headings) - 1):
        level, heading_line, start_i = headings[idx]
        next_level, _, next_start_i = headings[idx + 1]

        # Content is everything after the heading line until the next heading
        # i.e. lines[start_i+1 : next_start_i]
        block_content_lines = lines[start_i + 1 : next_start_i]
        block_content = "".join(block_content_lines).rstrip("\n")

        heading_blocks.append(
            {"level": level, "heading_text": heading_line, "content": block_content}
        )
    # (We skip the sentinel heading in the final result.)

    # We'll now convert heading_blocks into a tree using a stack-based approach
    root_nodes: List[Node] = []
    stack: List[Node] = []
    header_path: List[str] = []

    for hb in heading_blocks:
        level = hb["level"]  # type: ignore
        heading_txt = hb["heading_text"]
        content_txt = hb["content"]

        # --- Pop stack first! ---
        while stack and len(stack[-1].path) >= level:
            stack.pop()
            header_path.pop()

        # build new path, create a node for the heading
        new_path = header_path + [heading_txt]
        heading_node = Node(
            content=heading_txt, path=new_path, children=[]  # type: ignore
        )

        # Possibly create a content child for whatever lines were below the heading
        if content_txt.strip():  # type: ignore
            content_node = Node(
                content=content_txt, path=new_path, children=[]  # type: ignore
            )
            heading_node.children.append(content_node)

        # Attach heading_node to the stack top or as a root
        if stack:
            stack[-1].children.append(heading_node)
        else:
            root_nodes.append(heading_node)

        stack.append(heading_node)
        header_path.append(heading_txt)  # type: ignore

    return root_nodes


# The Chunk model for the final enriched chunks.
class Chunk(BaseModel):
    text: str  # The chunk text (which includes header context)
    path: List[str]  # The header path (list of header strings)
    token_count: int


# Configuration for chunking
class MarkdownChunkConfig(BaseModel):
    chunk_size: int = 200  # desired chunk size in tokens
    overlap_tokens: int = 30  # number of tokens to overlap between chunks
    variation_percent: float = 0.3  # allowed variation
    rollup: bool = True  # whether to roll up chunks
    header_context_sep: str = HEADER_CONTEXT_SEP  # separator for header context

    @field_validator("chunk_size", mode="before")
    @classmethod
    def convert_chunk_size_to_int(cls, v: Any) -> int:
        """Convert chunk_size to int, maintaining backward compatibility
        with Pydantic V1.
        """
        if isinstance(v, float):
            return int(v)
        return int(v)


# A simple tokenizer that counts tokens as whitespace-separated words.
def count_words(text: str) -> int:
    return len(text.split())


def recursive_chunk(text: str, config: MarkdownChunkConfig) -> List[str]:
    """
    Enhanced chunker that:
      1. Splits by paragraph (top-level).
      2. Splits paragraphs by sentences if needed (never mid-sentence unless huge).
      3. Allows going over the upper bound rather than splitting a single sentence.
      4. Overlaps only once between consecutive chunks.
      5. Looks ahead to avoid a "dangling" final chunk below the lower bound.
      6. Preserves \n\n (and other original spacing) as best as possible.
    """

    # -------------------------------------------------
    # Helpers
    # -------------------------------------------------
    def count_words(text_block: str) -> int:
        return len(text_block.split())

    lower_bound = int(config.chunk_size * (1 - config.variation_percent))
    upper_bound = int(config.chunk_size * (1 + config.variation_percent))

    # Quick check: if the entire text is short enough, return as-is.
    if count_words(text) <= upper_bound:
        return [text.strip()]

    # Split into paragraphs, preserving \n\n if it's there.
    raw_paragraphs = text.split("\n\n")
    paragraphs = []
    for i, p in enumerate(raw_paragraphs):
        if p.strip():
            # Re-append the double-newline if not the last piece
            if i < len(raw_paragraphs) - 1:
                paragraphs.append(p + "\n\n")
            else:
                paragraphs.append(p)

    # Split paragraphs into "segments": each segment is either
    # a full short paragraph or (if too big) a list of sentences.
    sentence_regex = r"(?<=[.!?])\s+"

    def split_paragraph_into_sentences(paragraph: str) -> List[str]:
        """
        Return a list of sentence-sized segments. If a single sentence
        is bigger than upper_bound, do a word-level fallback.
        """
        if count_words(paragraph) <= upper_bound:
            return [paragraph]

        sentences = re.split(sentence_regex, paragraph)
        # Clean up stray whitespace
        sentences = [s.strip() for s in sentences if s.strip()]

        expanded = []
        for s in sentences:
            if count_words(s) > upper_bound:
                expanded.extend(_fallback_word_split(s, config))
            else:
                expanded.append(s)
        return expanded

    def _fallback_word_split(long_text: str, cfg: MarkdownChunkConfig) -> List[str]:
        """
        As a last resort, split extremely large 'sentence' by words.
        """
        words = long_text.split()
        pieces = []
        start = 0
        while start < len(words):
            end = start + cfg.chunk_size
            chunk_words = words[start:end]
            pieces.append(" ".join(chunk_words))
            start = end
        return pieces

    # Build a list of segments
    segments = []
    for para in paragraphs:
        if count_words(para) > upper_bound:
            # split into sentences
            segs = split_paragraph_into_sentences(para)
            segments.extend(segs)
        else:
            segments.append(para)

    # -------------------------------------------------
    # Accumulate segments into final chunks
    # -------------------------------------------------
    chunks = []
    current_chunk = ""
    current_count = 0

    def flush_chunk() -> None:
        nonlocal current_chunk, current_count
        trimmed = current_chunk.strip()
        if trimmed:
            chunks.append(trimmed)
        current_chunk = ""
        current_count = 0

    def remaining_tokens_in_future(all_segments: List[str], current_index: int) -> int:
        """Sum of word counts from current_index onward."""
        return sum(count_words(s) for s in all_segments[current_index:])

    for i, seg in enumerate(segments):
        seg_count = count_words(seg)

        # If this single segment alone exceeds upper_bound, we accept it as a big chunk.
        if seg_count > upper_bound:
            # If we have something in the current chunk, flush it first
            flush_chunk()
            # Then store this large segment as its own chunk
            chunks.append(seg.strip())
            continue

        # Attempt to add seg to the current chunk
        if (current_count + seg_count) > upper_bound and (current_count >= lower_bound):
            # We would normally flush here, but let's see if we are nearing the end:
            # If the remaining tokens (including this one) is < lower_bound,
            # we just add it anyway to avoid creating a tiny final chunk.
            future_tokens = remaining_tokens_in_future(segments, i)
            if future_tokens < lower_bound:
                # Just add it (allowing to exceed upper bound)
                if current_chunk:
                    # Add space or preserve newline carefully
                    # We'll do a basic approach here:
                    if seg.startswith("\n\n"):
                        current_chunk += seg  # preserve double new line
                    else:
                        current_chunk += " " + seg
                    current_count = count_words(current_chunk)
                else:
                    current_chunk = seg
                    current_count = seg_count
            else:
                # Normal flush
                old_chunk = current_chunk
                flush_chunk()
                # Overlap from old_chunk
                overlap_tokens_list = (
                    old_chunk.split()[-config.overlap_tokens :] if old_chunk else []
                )
                overlap_str = (
                    " ".join(overlap_tokens_list) if overlap_tokens_list else ""
                )
                if overlap_str:
                    current_chunk = overlap_str + " " + seg
                else:
                    current_chunk = seg
                current_count = count_words(current_chunk)
        else:
            # Just accumulate
            if current_chunk:
                if seg.startswith("\n\n"):
                    current_chunk += seg
                else:
                    current_chunk += " " + seg
            else:
                current_chunk = seg
            current_count = count_words(current_chunk)

    # Flush leftover
    flush_chunk()

    # Return non-empty
    return [c for c in chunks if c.strip()]


# Function to process a Node and produce enriched chunks.
def chunk_node(node: Node, config: MarkdownChunkConfig) -> List[Chunk]:
    chunks: List[Chunk] = []

    # Check if this is a header-only node.
    is_header_only = node.path and node.content.strip() == node.path[-1]

    # Only generate a chunk for the node if it has non-header content,
    # or if it’s header-only AND has no children (i.e., it's a leaf header).
    if node.content.strip() and (not is_header_only or not node.children):
        header_prefix = (
            config.header_context_sep.join(node.path) + "\n\n" if node.path else ""
        )
        content_chunks = recursive_chunk(node.content, config)
        for chunk_text in content_chunks:
            full_text = header_prefix + chunk_text
            chunks.append(
                Chunk(
                    text=full_text, path=node.path, token_count=count_words(full_text)
                )
            )

    # Process children nodes recursively.
    for child in node.children:
        child_chunks = chunk_node(child, config)
        chunks.extend(child_chunks)

    return chunks


# Function to process an entire tree of Nodes.
def chunk_tree(root_nodes: List[Node], config: MarkdownChunkConfig) -> List[Chunk]:
    all_chunks: List[Chunk] = []
    for node in root_nodes:
        all_chunks.extend(chunk_node(node, config))
    return all_chunks


def aggregate_content(node: Node) -> str:
    """
    Recursively aggregate the content from a node and all its descendants,
    excluding header-only nodes to avoid duplication.
    """
    parts = []

    # Skip header-only nodes in content aggregation
    is_header_only = node.path and node.content.strip() == node.path[-1].strip()
    if not is_header_only and node.content.strip():
        parts.append(node.content.strip())

    # Recurse on children
    for child in node.children:
        child_text = aggregate_content(child)
        if child_text.strip():
            parts.append(child_text.strip())

    return "\n\n".join(parts)


def flatten_tree(node: Node, level: int = 0) -> str:
    """
    Flatten a node and its children back into proper markdown text.

    Args:
        node: The node to flatten
        level: The current heading level (depth in the tree)

    Returns:
        str: Properly formatted markdown text
    """
    result = ""

    # Check if this is a header node (content matches last item in path)
    is_header = node.path and node.content.strip().startswith("#")

    # For header nodes, don't duplicate the hash marks
    if is_header:
        result = node.content.strip() + "\n\n"
    elif node.content.strip():
        result = node.content.strip() + "\n\n"

    # Process all children
    for child in node.children:
        result += flatten_tree(child, level + 1)

    return result


def rollup_chunk_node(
    node: Node, config: MarkdownChunkConfig, prefix: str = ""
) -> List[Chunk]:
    """
    Recursively produce rollup chunks from `node`, passing down a `prefix`
    (e.g., parent heading(s)).

    - If a node is heading-only (content == last path item) and has children,
      we skip creating a chunk for that node alone and instead add that heading
      to the `prefix` for child nodes.
    - If a node is NOT heading-only OR has no children, we try to fit all of its
      flattened content into a single chunk. If it's too large, we chunk it.
    - We pass the (possibly updated) prefix down to children, so each child's
      chunk is enriched exactly once with all ancestor headings.
    """

    chunks: List[Chunk] = []

    # Check if the node is "heading-only" and has children
    # e.g. node.content=="# Chapter 1" and node.path[-1]=="# Chapter 1"
    is_heading_only_with_children = (
        node.path
        and node.content.strip() == node.path[-1].strip()
        and len(node.children) > 0
    )

    if is_heading_only_with_children:
        # We do NOT create a chunk for this node alone.
        # Instead, we add its heading to the prefix for child chunks.
        new_prefix = prefix + node.content.strip()
        for i, child in enumerate(node.children):
            sep = "\n\n" if i == 0 else config.header_context_sep
            chunks.extend(rollup_chunk_node(child, config, prefix=new_prefix + sep))
        return chunks

    # If not heading-only-with-children, we handle this node's own content:
    # Flatten the entire node (including sub-children) in standard Markdown form.
    flattened = flatten_tree(node, level=len(node.path))
    flattened_with_prefix = prefix + flattened
    total_tokens = count_words(flattened_with_prefix)

    # Check if we can roll up everything (node + children) in a single chunk
    if total_tokens <= config.chunk_size * (1 + config.variation_percent):
        # One single chunk for the entire subtree
        chunks.append(
            Chunk(text=flattened_with_prefix, path=node.path, token_count=total_tokens)
        )
    else:
        # It's too large overall. We'll chunk the node's own content first (if any),
        # then recurse on children.
        node_content = node.content.strip()

        # If we have actual content that is not just a heading, chunk it with the prefix
        # (like "preamble" text).
        # Note: if this node is heading-only but has NO children,
        # it will still land here
        # (because is_heading_only_with_children was False due to zero children).
        if node_content and (not node.path or node_content != node.path[-1].strip()):
            # The node is actual content (not purely heading).
            # We'll chunk it in paragraphs/sentences with the prefix.
            content_chunks = recursive_chunk(node_content, config)
            for text_block in content_chunks:
                block_with_prefix = prefix + text_block
                chunks.append(
                    Chunk(
                        text=block_with_prefix,
                        path=node.path,
                        token_count=count_words(block_with_prefix),
                    )
                )

        # Now recurse on children, passing the same prefix so they get it too
        for child in node.children:
            chunks.extend(rollup_chunk_node(child, config, prefix=prefix))

    return chunks


def rollup_chunk_tree(
    root_nodes: List[Node],
    config: MarkdownChunkConfig,
) -> List[Chunk]:
    # Create a dummy root node that contains everything.
    dummy_root = Node(content="", path=[], children=root_nodes)

    # Now process just the dummy root node with an empty prefix.
    chunks = rollup_chunk_node(dummy_root, config, prefix="")
    return chunks


def chunk_markdown(markdown_text: str, config: MarkdownChunkConfig) -> List[str]:
    tree = parse_markdown_headings(markdown_text)
    if len(tree) == 1 and len(tree[0].children) == 0:
        # Pure text, no hierarchy, so just use recursive_chunk
        text_chunks = recursive_chunk(markdown_text, config)
        return [_cleanup_text(chunk) for chunk in text_chunks]
    if config.rollup:
        chunks = rollup_chunk_tree(tree, config)
    else:
        chunks = chunk_tree(tree, config)
    return [_cleanup_text(chunk.text) for chunk in chunks]


if __name__ == "__main__":
    # Example usage:
    markdown_text = """# Title
Intro para. Hope this is not
getting split.
## SubTitle
- Item1
- Item2
"""
    # Set up chunking config with very large chunk size.
    # (you can adjust chunk_size, overlap_tokens, variation_percent)
    config = MarkdownChunkConfig(
        chunk_size=200, overlap_tokens=5, variation_percent=0.2
    )
    chunks = chunk_markdown(markdown_text, config)

    for idx, chunk in enumerate(chunks, 1):
        print(f"--- Chunk {idx} --- ")
        print(chunk)
        print()

    config.rollup = True
    # with rollup_chunk_tree we get entire doc as 1 chunk
    chunks = chunk_markdown(markdown_text, config)
    assert len(chunks) == 1
    for idx, chunk in enumerate(chunks, 1):
        print(f"--- Chunk {idx} ---")
        print(chunk)
        print()
</file>

<file path="langroid/parsing/para_sentence_split.py">
import re
from typing import Callable, List

from bs4 import BeautifulSoup


def remove_extra_whitespace(s: str) -> str:
    lines = s.split("\n")
    cleaned_lines = [" ".join(line.split()) for line in lines]
    return "\n".join(cleaned_lines)


def custom_sent_tokenize(text: str) -> List[str]:
    sentences = [
        sentence.strip()
        for sentence in re.split(r"\.\s|\.\n", text)
        if sentence.strip()
    ]
    # append a period if the sentence does not end with one
    return [s + "." if s[-1] != "." else s for s in sentences]


def create_chunks(
    text: str, chunk_size: int, length_fn: Callable[[str], int]
) -> List[str]:
    def _chunk_sentences(sentences: List[str], chunk_size: int) -> List[str]:
        chunks = []
        current_chunk: List[str] = []
        current_chunk_length = 0

        for sentence in sentences:
            sentence_length = length_fn(sentence)
            if current_chunk_length + sentence_length > chunk_size:
                if current_chunk:
                    chunks.append(" ".join(current_chunk))
                current_chunk = [sentence]
                current_chunk_length = sentence_length
            else:
                current_chunk.append(sentence)
                current_chunk_length += sentence_length

        if current_chunk:
            new_chunk = " ".join(current_chunk).strip()
            if new_chunk:
                chunks.append(" ".join(current_chunk).strip())

        return chunks

    soup = BeautifulSoup(text, "html.parser")
    text = soup.get_text()
    # First, try to split the document into paragraphs
    paragraphs = text.split("\n\n")

    # If paragraphs are too long, split them into sentences
    if any(length_fn(p) > chunk_size for p in paragraphs):
        sentences = custom_sent_tokenize(text)
        chunks = _chunk_sentences(sentences, chunk_size)
    else:
        chunks = paragraphs

    chunks = [chunk.strip() for chunk in chunks if chunk.strip() != ""]
    return chunks
</file>

<file path="langroid/parsing/parse_json.py">
import ast
import json
from datetime import datetime
from typing import Any, Dict, Iterator, List, Union

import yaml
from json_repair import repair_json
from pyparsing import nested_expr, original_text_for


def is_valid_json(json_str: str) -> bool:
    """Check if the input string is a valid JSON.

    Args:
        json_str (str): The input string to check.

    Returns:
        bool: True if the input string is a valid JSON, False otherwise.
    """
    try:
        json.loads(json_str)
        return True
    except ValueError:
        return False


def flatten(nested_list) -> Iterator[str]:  # type: ignore
    """Flatten a nested list into a single list of strings"""
    for item in nested_list:
        if isinstance(item, (list, tuple)):
            for subitem in flatten(item):
                yield subitem
        else:
            yield item


def get_json_candidates(s: str) -> List[str]:
    """Get top-level JSON candidates, i.e. strings between curly braces."""
    # Define the grammar for matching curly braces
    curly_braces = original_text_for(nested_expr("{", "}"))

    # Parse the string
    try:
        results = curly_braces.search_string(s)
        # Properly convert nested lists to strings
        return [r[0] for r in results]
    except Exception:
        return []


def parse_imperfect_json(json_string: str) -> Union[Dict[str, Any], List[Any]]:
    if not json_string.strip():
        raise ValueError("Empty string is not valid JSON")

    # First, try parsing with ast.literal_eval
    try:
        result = ast.literal_eval(json_string)
        if isinstance(result, (dict, list)):
            return result
    except (ValueError, SyntaxError):
        pass

    # If ast.literal_eval fails or returns non-dict/list, try repair_json
    json_repaired_obj = repair_json(json_string, return_objects=True)
    if isinstance(json_repaired_obj, (dict, list)):
        return json_repaired_obj
    else:
        try:
            # fallback on yaml
            yaml_result = yaml.safe_load(json_string)
            if isinstance(yaml_result, (dict, list)):
                return yaml_result
        except yaml.YAMLError:
            pass

    # If all methods fail, raise ValueError
    raise ValueError(f"Unable to parse as JSON: {json_string}")


def try_repair_json_yaml(s: str) -> str | None:
    """
    Attempt to load as json, and if it fails, try repairing the JSON.
    If that fails, replace any \n with space as a last resort.
    NOTE - replacing \n with space will result in format loss,
    which may matter in generated code (e.g. python, toml, etc)
    """
    s_repaired_obj = repair_json(s, return_objects=True)
    if isinstance(s_repaired_obj, list):
        if len(s_repaired_obj) > 0:
            s_repaired_obj = s_repaired_obj[0]
        else:
            s_repaired_obj = None
    if s_repaired_obj is not None:
        return json.dumps(s_repaired_obj)  # type: ignore
    else:
        try:
            yaml_result = yaml.safe_load(s)
            if isinstance(yaml_result, dict):
                return json.dumps(yaml_result)
        except yaml.YAMLError:
            pass
        # If it still fails, replace any \n with space as a last resort
        s = s.replace("\n", " ")
        if is_valid_json(s):
            return s
        else:
            return None  # all failed


def extract_top_level_json(s: str) -> List[str]:
    """Extract all top-level JSON-formatted substrings from a given string.

    Args:
        s (str): The input string to search for JSON substrings.

    Returns:
        List[str]: A list of top-level JSON-formatted substrings.
    """
    # Find JSON object and array candidates
    json_candidates = get_json_candidates(s)
    maybe_repaired_jsons = map(try_repair_json_yaml, json_candidates)

    return [candidate for candidate in maybe_repaired_jsons if candidate is not None]


def top_level_json_field(s: str, f: str) -> Any:
    """
    Extract the value of a field f from a top-level JSON object.
    If there are multiple, just return the first.

    Args:
        s (str): The input string to search for JSON substrings.
        f (str): The field to extract from the JSON object.

    Returns:
        str: The value of the field f in the top-level JSON object, if any.
            Otherwise, return an empty string.

    Note:
        This function is designed to never crash. If any exception occurs during
        JSON parsing or field extraction, it gracefully returns an empty string.
    """
    try:
        jsons = extract_top_level_json(s)
        if len(jsons) == 0:
            return ""
        for j in jsons:
            try:
                json_data = json.loads(j)
                if isinstance(json_data, dict):
                    if f in json_data:
                        return json_data[f]
                elif isinstance(json_data, list):
                    # Some responses wrap candidate JSON objects in a list; scan them.
                    for item in json_data:
                        if isinstance(item, dict) and f in item:
                            return item[f]
            except (json.JSONDecodeError, TypeError, KeyError):
                # If this specific JSON fails to parse, continue to next candidate
                continue
    except Exception:
        # Catch any unexpected errors to ensure we never crash
        pass

    return ""


def datetime_to_json(obj: Any) -> Any:
    if isinstance(obj, datetime):
        return obj.isoformat()
    # Let json.dumps() handle the raising of TypeError for non-serializable objects
    return obj
</file>

<file path="langroid/parsing/pdf_utils.py">
import tempfile
from io import BytesIO
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TYPE_CHECKING, Any, BinaryIO, List, Optional, Tuple, Union

try:
    import fitz
except ImportError:
    if not TYPE_CHECKING:
        fitz = None

from langroid.exceptions import LangroidImportError

if fitz is None:
    raise LangroidImportError("fitz", ["pymupdf", "all", "pdf-parsers", "doc-chat"])


def pdf_split_pages(
    input_pdf: Union[BytesIO, BinaryIO, str],
    splits: Optional[List[int]] = None,
) -> Tuple[List[Path], TemporaryDirectory[Any]]:
    """Splits a PDF into individual pages or chunks in a temporary directory.

    Args:
        input_pdf: Input PDF file in bytes, binary mode, or a file path
        splits: Optional list of page numbers to split at.
                If provided, pages will be grouped into chunks ending at
                these page numbers.
                For example, if splits = [4, 9], the result will have pages 1-4, 5-9,
                and 10-end.
                If not provided, default to splitting into individual pages.
        max_workers: Maximum number of concurrent workers for parallel processing

    Returns:
        Tuple containing:
            - List of paths to individual PDF pages or chunks
            - Temporary directory object (caller must call cleanup())

    Example:
        paths, tmp_dir = split_pdf_temp("input.pdf")
        # Use paths...
        tmp_dir.cleanup()  # Clean up temp files when done
    """
    tmp_dir = tempfile.TemporaryDirectory()
    if isinstance(input_pdf, str):
        doc = fitz.open(input_pdf)
    else:
        doc = fitz.open(stream=input_pdf, filetype="pdf")
    paths = []

    total_pages = len(doc)

    if splits is None:
        # Split into individual pages (original behavior)
        for page_num in range(total_pages):
            new_doc = fitz.open()
            new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)
            output = Path(tmp_dir.name) / f"page_{page_num + 1}.pdf"
            new_doc.save(str(output))
            new_doc.close()
            paths.append(output)
    else:
        # Split according to specified page ranges
        # Make sure the splits list is sorted and includes all valid splits
        splits = sorted([s for s in splits if 1 <= s <= total_pages])

        # Create the ranges to process
        ranges = []
        start_page = 0
        for end_page in splits:
            ranges.append((start_page, end_page - 1))
            start_page = end_page

        # Add the final range if there are pages after the last split
        if start_page < total_pages:
            ranges.append((start_page, total_pages - 1))

        # Process each range
        for i, (from_page, to_page) in enumerate(ranges):
            new_doc = fitz.open()
            new_doc.insert_pdf(doc, from_page=from_page, to_page=to_page)
            output = Path(tmp_dir.name) / f"pages_{from_page + 1}_to_{to_page + 1}.pdf"
            new_doc.save(str(output))
            new_doc.close()
            paths.append(output)

    doc.close()
    return paths, tmp_dir
</file>

<file path="langroid/parsing/routing.py">
import re
from typing import Optional, Tuple


def parse_addressed_message(
    content: str, addressing: str = "@"
) -> Tuple[Optional[str], str]:
    """In a message-string containing possibly multiple @<recipient> occurrences,
    find the last addressee and extract their name,
    and the message content following it.

    E.g. "thank you @bob, now I will ask @alice again. @alice, where is the mirror?" =>
    ("alice", "where is the mirror?")

    Args:
        content (str): The message content.
        addressing (str, optional): The addressing character. Defaults to "@".

    Returns:
        Tuple[Optional[str], str]:
        A tuple containing the last addressee and the subsequent message content.
    """
    # Regex to find all occurrences of the pattern
    pattern = re.compile(rf"{re.escape(addressing)}(\w+)[^\w]")
    matches = list(pattern.finditer(content))

    if not matches:
        return None, content  # No addressee found, return None and original content

    # Get the last match
    last_match = matches[-1]
    last_addressee = last_match.group(1)
    # Extract content after the last addressee
    content_after = content[last_match.end() :].strip()

    return last_addressee, content_after
</file>

<file path="langroid/parsing/search.py">
"""
Utils to search for close matches in (a list of) strings.
Useful for retrieval of docs/chunks relevant to a query, in the context of
Retrieval-Augmented Generation (RAG), and SQLChat (e.g., to pull relevant parts of a
large schema).
See tests for examples: tests/main/test_string_search.py
"""

import difflib
import re
from typing import List, Tuple

from rank_bm25 import BM25Okapi
from thefuzz import fuzz, process

from langroid.mytypes import Document

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,
) -> List[Tuple[Document, float]]:
    """
    Find approximate matches of the query in the docs and return surrounding
    characters.

    Args:
        query (str): The search string.
        docs (List[Document]): List of Document objects to search through.
        docs_clean (List[Document]): List of Document objects with cleaned content.
        k (int): Number of best matches to return.
        words_before (int|None): Number of words to include before each match.
            Default None => return max
        words_after (int|None): Number of words to include after each match.
            Default None => return max

    Returns:
        List[Tuple[Document,float]]: List of (Document, score) tuples.
    """
    if len(docs) == 0:
        return []
    best_matches = process.extract(
        query,
        [d.content for d in docs_clean],
        limit=k,
        scorer=fuzz.partial_ratio,
    )

    real_matches = [(m, score) for m, score in best_matches if score > 50]
    # find the original docs that corresponding to the matches
    orig_doc_matches = []
    for i, (m, s) in enumerate(real_matches):
        for j, doc_clean in enumerate(docs_clean):
            if m in doc_clean.content:
                orig_doc_matches.append((docs[j], s))
                break
    if words_after is None and words_before is None:
        return orig_doc_matches
    if len(orig_doc_matches) == 0:
        return []
    if set(orig_doc_matches[0][0].model_fields) != {"content", "metadata"}:
        # If there are fields beyond just content and metadata,
        # we do NOT want to create new document objects with content fields
        # based on words_before and words_after, since we don't know how to
        # set those other fields.
        return orig_doc_matches

    contextual_matches = []
    for match, score 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,
                    ),
                    score,
                )
            )

    return contextual_matches


def preprocess_text(text: str) -> str:
    """
    Preprocesses the given text by:
    1. Lowercasing all words.
    2. Tokenizing (splitting the text into words).
    3. Removing punctuation.
    4. Removing stopwords.
    5. Lemmatizing words.

    Args:
        text (str): The input text.

    Returns:
        str: The preprocessed text.
    """
    # Ensure the NLTK resources are available
    for resource in ["tokenizers/punkt", "corpora/wordnet", "corpora/stopwords"]:
        download_nltk_resource(resource)
    from nltk.corpus import stopwords
    from nltk.stem import WordNetLemmatizer
    from nltk.tokenize import RegexpTokenizer

    # Lowercase the text
    text = text.lower()

    # Tokenize the text and remove punctuation
    tokenizer = RegexpTokenizer(r"\w+")
    tokens = tokenizer.tokenize(text)

    # Remove stopwords
    stop_words = set(stopwords.words("english"))
    tokens = [t for t in tokens if t not in stop_words]

    # Lemmatize words
    lemmatizer = WordNetLemmatizer()
    tokens = [lemmatizer.lemmatize(t) for t in tokens]

    # Join the words back into a string
    text = " ".join(tokens)

    return text


def find_closest_matches_with_bm25(
    docs: List[Document],
    docs_clean: List[Document],
    query: str,
    k: int = 5,
) -> List[Tuple[Document, float]]:
    """
    Finds the k closest approximate matches using the BM25 algorithm.

    Args:
        docs (List[Document]): List of Documents to search through.
        docs_clean (List[Document]): List of cleaned Documents
        query (str): The search query.
        k (int, optional): Number of matches to retrieve. Defaults to 5.

    Returns:
        List[Tuple[Document,float]]: List of (Document, score) tuples.
    """
    if len(docs) == 0:
        return []
    texts = [doc.content for doc in docs_clean]
    query = preprocess_text(query)

    text_words = [text.split() for text in texts]

    bm25 = BM25Okapi(text_words)
    query_words = query.split()
    doc_scores = bm25.get_scores(query_words)

    # Get indices of top k scores
    top_indices = sorted(range(len(doc_scores)), key=lambda i: -doc_scores[i])[:k]

    # return the original docs, based on the scores from cleaned docs
    return [(docs[i], doc_scores[i]) for i in top_indices]


def get_context(
    query: str,
    text: str,
    words_before: int | None = 100,
    words_after: int | None = 100,
) -> 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.

    Args:
    query (str): The string to search for.
    text (str): The body of text in which to search.
    b (int): The number of words before the query to return.
    a (int): The number of words after the query to return.

    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.
        The text is extracted from the original `text`, preserving formatting,
        whitespace, etc, so it does not disturb any downstream processing.
        If no 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)
    # 'fox jumps over the apple.'
    """

    # If no word limits specified, return full text
    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, 0, 0

    # make sure there is a good enough match to the query
    if fuzz.partial_ratio(query, text) < 40:
        return "", 0, 0

    # Find best matching position of query in text
    sequence_matcher = difflib.SequenceMatcher(None, text, query)
    match = sequence_matcher.find_longest_match(0, len(text), 0, len(query))

    if match.size == 0:
        return "", 0, 0

    # Count words before match point
    segments = text.split()
    n_segs = len(segments)
    start_segment_pos = len(text[: match.a].split())

    # Calculate word window boundaries
    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()))

    # Find character positions where words start
    word_positions = [m.start() for m in re.finditer(r"\S+", text)]

    # Convert word positions to character positions
    start_char = word_positions[start_pos] if start_pos < len(word_positions) else 0
    end_char = word_positions[min(end_pos, len(word_positions) - 1)] + len(
        text.split()[min(end_pos - 1, len(word_positions) - 1)]
    )

    # return exact substring with original formatting
    return text[start_char:end_char], start_pos, end_pos


def eliminate_near_duplicates(passages: List[str], threshold: float = 0.8) -> List[str]:
    """
    Eliminate near duplicate text passages from a given list using MinHash and LSH.
    TODO: this has not been tested and the datasketch lib is not a dependency.
    Args:
        passages (List[str]): A list of text passages.
        threshold (float, optional): Jaccard similarity threshold to consider two
                                     passages as near-duplicates. Default is 0.8.

    Returns:
        List[str]: A list of passages after eliminating near duplicates.

    Example:
        passages = ["Hello world", "Hello, world!", "Hi there", "Hello world!"]
        print(eliminate_near_duplicates(passages))
        # ['Hello world', 'Hi there']
    """

    from datasketch import MinHash, MinHashLSH

    # Create LSH index
    lsh = MinHashLSH(threshold=threshold, num_perm=128)

    # Create MinHash objects for each passage and insert to LSH
    minhashes = {}
    for idx, passage in enumerate(passages):
        m = MinHash(num_perm=128)
        for word in passage.split():
            m.update(word.encode("utf-8"))
        lsh.insert(idx, m)
        minhashes[idx] = m

    unique_idxs = set()
    for idx in minhashes.keys():
        # Query for similar passages (including itself)
        result = lsh.query(minhashes[idx])

        # If only the passage itself is returned, it's unique
        if len(result) == 1 and idx in result:
            unique_idxs.add(idx)

    return [passages[idx] for idx in unique_idxs]
</file>

<file path="langroid/parsing/spider.py">
from typing import List, Set, no_type_check
from urllib.parse import urlparse

from langroid.exceptions import LangroidImportError

try:
    from pydispatch import dispatcher
    from scrapy import signals
    from scrapy.crawler import CrawlerRunner
    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")


@no_type_check
class DomainSpecificSpider(CrawlSpider):  # type: ignore
    name = "domain_specific_spider"

    custom_settings = {"DEPTH_LIMIT": 1, "CLOSESPIDER_ITEMCOUNT": 20}

    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.

        Args:
            start_url (str): The starting URL.
            k (int, optional): The max desired final URLs. Defaults to 20.
        """
        super(DomainSpecificSpider, self).__init__(*args, **kwargs)
        self.start_urls = [start_url]
        self.allowed_domains = [urlparse(start_url).netloc]
        self.k = k
        self.visited_urls: Set[str] = set()

    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 LxmlLinkExtractor(allow_domains=self.allowed_domains).extract_links(
            response
        ):
            if len(self.visited_urls) < self.k:
                self.visited_urls.add(link.url)
                yield {"url": link.url}


@no_type_check
def scrapy_fetch_urls(url: str, k: int = 20) -> List[str]:
    """Fetches up to k URLs reachable from the input URL using Scrapy.

    Args:
        url (str): The starting URL.
        k (int, optional): The max desired final URLs. Defaults to 20.

    Returns:
        List[str]: List of URLs within the same domain as the input URL.
    """
    urls = []

    def _collect_urls(spider):
        """Handler for the spider_closed signal. Collects the visited URLs."""
        nonlocal urls
        urls.extend(list(spider.visited_urls))

    # Connect the spider_closed signal with our handler
    dispatcher.connect(_collect_urls, signal=signals.spider_closed)

    runner = CrawlerRunner(
        {
            "USER_AGENT": "Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)"
        }
    )

    d = runner.crawl(DomainSpecificSpider, start_url=url, k=k)

    # Block until crawling is done and then stop the reactor
    crawl_deferred = defer.Deferred()

    def _crawl_done(_):
        reactor.stop()
        crawl_deferred.callback(urls)

    d.addBoth(_crawl_done)

    # Start the reactor, it will stop once the crawl is done
    reactor.run(installSignalHandlers=0)

    # This will block until the deferred gets a result
    return crawl_deferred.result


# Test the function
if __name__ == "__main__":
    fetched_urls = scrapy_fetch_urls("https://example.com", 5)
    for url in fetched_urls:
        print(url)
</file>

<file path="langroid/parsing/table_loader.py">
from csv import Sniffer
from typing import List

import pandas as pd


def read_tabular_data(path_or_url: str, sep: None | str = None) -> pd.DataFrame:
    """
    Reads tabular data from a file or URL and returns a pandas DataFrame.
    The separator is auto-detected if not specified.

    Args:
        path_or_url (str): Path or URL to the file to be read.

    Returns:
        pd.DataFrame: Data from file or URL as a pandas DataFrame.

    Raises:
        ValueError: If the data cannot be read or is misformatted.
    """
    try:
        if sep is None:
            # Read the first few lines to guess the separator
            with pd.io.common.get_handle(path_or_url, "r") as file_handler:
                first_lines = "".join(file_handler.handle.readlines(5))
                sep = Sniffer().sniff(first_lines).delimiter
                # If it's a local file, reset to the beginning
                if hasattr(file_handler.handle, "seek"):
                    file_handler.handle.seek(0)

        # Read the data

        # get non-blank column names
        with pd.io.common.get_handle(path_or_url, "r") as f:
            header_line = f.handle.readline().strip()
            valid_cols = [col for col in header_line.split(sep) if col]
            valid_cols = [c.replace('"', "").replace("'", "") for c in valid_cols]
            if hasattr(f.handle, "seek"):
                f.handle.seek(0)

        # use only those columns
        data = pd.read_csv(path_or_url, sep=sep, usecols=valid_cols)
        data.columns = data.columns.str.strip()  # e.g. "  column 1  " -> "column 1"

        return data

    except Exception as e:
        raise ValueError(
            "Unable to read data. "
            "Please ensure it is correctly formatted. Error: " + str(e)
        )


def describe_dataframe(
    df: pd.DataFrame, filter_fields: List[str] = [], n_vals: int = 10
) -> str:
    """
    Generates a description of the columns in the dataframe,
    along with a listing of up to `n_vals` unique values for each column.
    Intended to be used to insert into an LLM context so it can generate
    appropriate queries or filters on the df.

    Args:
    df (pd.DataFrame): The dataframe to describe.
    filter_fields (list): A list of fields that can be used for filtering.
        When non-empty, the values-list will be restricted to these.
    n_vals (int): How many unique values to show for each column.

    Returns:
    str: A description of the dataframe.
    """
    description = []
    for column in df.columns.to_list():
        unique_values = df[column].dropna().unique()
        unique_count = len(unique_values)
        if column not in filter_fields:
            values_desc = f"{unique_count} unique values"
        else:
            if unique_count > n_vals:
                displayed_values = unique_values[:n_vals]
                more_count = unique_count - n_vals
                values_desc = f" Values - {displayed_values}, ... {more_count} more"
            else:
                values_desc = f" Values - {unique_values}"
        col_type = "string" if df[column].dtype == "object" else df[column].dtype
        col_desc = f"* {column} ({col_type}); {values_desc}"
        description.append(col_desc)

    all_cols = "\n".join(description)

    return f"""
        Name of each field, its type and unique values (up to {n_vals}):
        {all_cols}
        """
</file>

<file path="langroid/parsing/url_loader.py">
import asyncio
import logging
import os
from abc import ABC, abstractmethod
from tempfile import NamedTemporaryFile
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional

import markdownify as md
from dotenv import load_dotenv
from pydantic_settings import BaseSettings, SettingsConfigDict

from langroid.exceptions import LangroidImportError
from langroid.mytypes import DocMetaData, Document
from langroid.parsing.document_parser import DocumentParser, ImagePdfParser
from langroid.parsing.parser import Parser, ParsingConfig

if TYPE_CHECKING:
    from firecrawl import FirecrawlApp

    try:
        from crawl4ai import CrawlResult
        from crawl4ai.async_configs import BrowserConfig, CrawlerRunConfig
        from crawl4ai.content_scraping_strategy import ContentScrapingStrategy
        from crawl4ai.deep_crawling import DeepCrawlStrategy
        from crawl4ai.extraction_strategy import ExtractionStrategy
        from crawl4ai.markdown_generation_strategy import MarkdownGenerationStrategy
    except ImportError:
        raise LangroidImportError("crawl4ai", "crawl4ai")

load_dotenv()

logging.getLogger("url_loader").setLevel(logging.WARNING)


# Base crawler config and specific configurations
class BaseCrawlerConfig(BaseSettings):
    """Base configuration for web crawlers."""

    parser: Optional[Parser] = None


class TrafilaturaConfig(BaseCrawlerConfig):
    """Configuration for Trafilatura crawler."""

    threads: int = 4
    format: str = "markdown"  # or "xml" or "txt"


class FirecrawlConfig(BaseCrawlerConfig):
    """Configuration for Firecrawl crawler."""

    api_key: str = ""
    mode: str = "scrape"
    params: Dict[str, Any] = {}
    timeout: Optional[int] = None

    model_config = SettingsConfigDict(env_prefix="FIRECRAWL_")


class ExaCrawlerConfig(BaseCrawlerConfig):
    api_key: str = ""

    model_config = SettingsConfigDict(env_prefix="EXA_")


class Crawl4aiConfig(BaseCrawlerConfig):
    """Configuration for the Crawl4aiCrawler."""

    crawl_mode: Literal["simple", "deep"] = "simple"
    extraction_strategy: Optional["ExtractionStrategy"] = None
    markdown_strategy: Optional["MarkdownGenerationStrategy"] = None
    deep_crawl_strategy: Optional["DeepCrawlStrategy"] = None
    scraping_strategy: Optional["ContentScrapingStrategy"] = None
    browser_config: Optional["BrowserConfig"] = None
    run_config: Optional["CrawlerRunConfig"] = None

    model_config = SettingsConfigDict(arbitrary_types_allowed=True)


# Resolve forward references for Crawl4aiConfig after the class is defined
try:
    from crawl4ai.async_configs import BrowserConfig, CrawlerRunConfig
    from crawl4ai.content_scraping_strategy import ContentScrapingStrategy
    from crawl4ai.deep_crawling import DeepCrawlStrategy
    from crawl4ai.extraction_strategy import ExtractionStrategy
    from crawl4ai.markdown_generation_strategy import MarkdownGenerationStrategy

    # Rebuild the model with resolved references
    Crawl4aiConfig.model_rebuild()
except ImportError:
    # If crawl4ai is not installed, leave forward refs as strings
    pass


class BaseCrawler(ABC):
    """Abstract base class for web crawlers."""

    def __init__(self, config: BaseCrawlerConfig):
        """Initialize the base crawler.

        Args:
            config: Configuration for the crawler
        """
        self.parser = config.parser if self.needs_parser else None
        self.config: BaseCrawlerConfig = config

    @property
    @abstractmethod
    def needs_parser(self) -> bool:
        """Indicates whether the crawler requires a parser."""
        pass

    @abstractmethod
    def crawl(self, urls: List[str]) -> List[Document]:
        pass

    def _process_document(self, url: str) -> List[Document]:
        if self.parser:
            import requests
            from requests.structures import CaseInsensitiveDict

            if self._is_document_url(url):
                try:
                    doc_parser = DocumentParser.create(url, self.parser.config)
                    new_chunks = doc_parser.get_doc_chunks()
                    if not new_chunks:
                        # If the document is empty, try to extract images
                        img_parser = ImagePdfParser(url, self.parser.config)
                        new_chunks = img_parser.get_doc_chunks()
                    return new_chunks
                except Exception as e:
                    logging.error(f"Error parsing {url}: {e}")
                    return []

            else:
                try:
                    headers = requests.head(url).headers
                except Exception as e:
                    logging.warning(f"Error getting headers for {url}: {e}")
                    headers = CaseInsensitiveDict()

                content_type = headers.get("Content-Type", "").lower()
                temp_file_suffix = None
                if "application/pdf" in content_type:
                    temp_file_suffix = ".pdf"
                elif (
                    "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
                    in content_type
                ):
                    temp_file_suffix = ".docx"
                elif "application/msword" in content_type:
                    temp_file_suffix = ".doc"

                if temp_file_suffix:
                    try:
                        response = requests.get(url)
                        with NamedTemporaryFile(
                            delete=False, suffix=temp_file_suffix
                        ) as temp_file:
                            temp_file.write(response.content)
                            temp_file_path = temp_file.name
                        doc_parser = DocumentParser.create(
                            temp_file_path, self.parser.config
                        )
                        docs = doc_parser.get_doc_chunks()
                        os.remove(temp_file_path)
                        return docs
                    except Exception as e:
                        logging.error(f"Error downloading/parsing {url}: {e}")
                        return []
        return []

    def _is_document_url(self, url: str) -> bool:
        return any(url.lower().endswith(ext) for ext in [".pdf", ".docx", ".doc"])


class CrawlerFactory:
    """Factory for creating web crawlers."""

    @staticmethod
    def create_crawler(config: BaseCrawlerConfig) -> BaseCrawler:
        """Create a crawler instance based on configuration type.

        Args:
            config: Configuration for the crawler

        Returns:
            A BaseCrawler instance

        Raises:
            ValueError: If config type is not supported
        """
        if isinstance(config, TrafilaturaConfig):
            return TrafilaturaCrawler(config)
        elif isinstance(config, FirecrawlConfig):
            return FirecrawlCrawler(config)
        elif isinstance(config, ExaCrawlerConfig):
            return ExaCrawler(config)
        elif isinstance(config, Crawl4aiConfig):
            return Crawl4aiCrawler(config)
        else:
            raise ValueError(f"Unsupported crawler configuration type: {type(config)}")


class TrafilaturaCrawler(BaseCrawler):
    """Crawler implementation using Trafilatura."""

    def __init__(self, config: TrafilaturaConfig):
        """Initialize the Trafilatura crawler.

        Args:
            config: Configuration for the crawler
        """
        super().__init__(config)
        self.config: TrafilaturaConfig = config

    @property
    def needs_parser(self) -> bool:
        return True

    def crawl(self, urls: List[str]) -> List[Document]:
        import trafilatura
        from trafilatura.downloads import (
            add_to_compressed_dict,
            buffered_downloads,
            load_download_buffer,
        )

        docs = []
        dl_dict = add_to_compressed_dict(urls)

        while not dl_dict.done:
            buffer, dl_dict = load_download_buffer(dl_dict, sleep_time=5)
            for url, result in buffered_downloads(buffer, self.config.threads):
                parsed_doc = self._process_document(url)
                if parsed_doc:
                    docs.extend(parsed_doc)
                else:
                    text = trafilatura.extract(
                        result,
                        no_fallback=False,
                        favor_recall=True,
                        include_formatting=True,
                        output_format=self.config.format,
                        with_metadata=True,  # Title, date, author... at start of text
                    )
                    if self.config.format in ["xml", "html"]:
                        # heading_style="ATX" for markdown headings, i.e. #, ##, etc.
                        text = md.markdownify(text, heading_style="ATX")
                    if text is None and result is not None and isinstance(result, str):
                        text = result
                    if text:
                        docs.append(
                            Document(content=text, metadata=DocMetaData(source=url))
                        )

        return docs


class FirecrawlCrawler(BaseCrawler):
    """Crawler implementation using Firecrawl."""

    def __init__(self, config: FirecrawlConfig) -> None:
        """Initialize the Firecrawl crawler.

        Args:
            config: Configuration for the crawler
        """
        super().__init__(config)
        self.config: FirecrawlConfig = config

    @property
    def needs_parser(self) -> bool:
        return False

    def _return_save_incremental_results(
        self, app: "FirecrawlApp", crawl_id: str, output_dir: str = "firecrawl_output"
    ) -> List[Document]:
        # Code used verbatim from firecrawl blog with few modifications
        # https://www.firecrawl.dev/blog/mastering-the-crawl-endpoint-in-firecrawl
        import json
        import time
        from pathlib import Path

        from tqdm import tqdm

        pbar = tqdm(desc="Pages saved", unit=" pages", dynamic_ncols=True)
        Path(output_dir).mkdir(parents=True, exist_ok=True)
        processed_urls: set[str] = set()
        docs = []

        while True:
            # Check current status
            status = app.check_crawl_status(crawl_id)
            new_pages = 0

            # Save new pages
            for page in status["data"]:
                url = page["metadata"]["url"]
                if url not in processed_urls:
                    content = page.get("markdown", "")
                    filename = f"{output_dir}/{len(processed_urls)}.md"
                    with open(filename, "w") as f:
                        f.write(content)
                    docs.append(
                        Document(
                            content=content,
                            metadata=DocMetaData(
                                source=url,
                                title=page["metadata"].get("title", "Unknown Title"),
                            ),
                        )
                    )
                    processed_urls.add(url)
                    new_pages += 1
            pbar.model_copy(update=new_pages)  # Update progress bar with new pages

            # Break if crawl is complete
            if status["status"] == "completed":
                print(f"Saved {len(processed_urls)} pages.")
                with open(f"{output_dir}/full_results.json", "w") as f:
                    json.dump(status, f, indent=2)
                break

            time.sleep(5)  # Wait before checking again
        return docs

    def crawl(self, urls: List[str]) -> List[Document]:
        try:
            from firecrawl import FirecrawlApp
        except ImportError:
            raise LangroidImportError("firecrawl", "firecrawl")

        app = FirecrawlApp(api_key=self.config.api_key)
        docs = []
        params = self.config.params.copy()  # Create a copy of the existing params

        if self.config.timeout is not None:
            params["timeout"] = self.config.timeout  # Add/override timeout in params

        if self.config.mode == "scrape":
            for url in urls:
                try:
                    result = app.scrape_url(url, params=params)
                    metadata = result.get(
                        "metadata", {}
                    )  # Default to empty dict if missing
                    status_code = metadata.get("statusCode")

                    if status_code == 200:
                        docs.append(
                            Document(
                                content=result["markdown"],
                                metadata=DocMetaData(
                                    source=url,
                                    title=metadata.get("title", "Unknown Title"),
                                ),
                            )
                        )
                except Exception as e:
                    logging.warning(
                        f"Firecrawl encountered an error for {url}: {e}. "
                        "Skipping but continuing."
                    )
        elif self.config.mode == "crawl":
            if not isinstance(urls, list) or len(urls) != 1:
                raise ValueError(
                    "Crawl mode expects 'urls' to be a list containing a single URL."
                )

            # Start the crawl
            crawl_status = app.async_crawl_url(url=urls[0], params=params)

            # Save results incrementally
            docs = self._return_save_incremental_results(app, crawl_status["id"])
        return docs


class ExaCrawler(BaseCrawler):
    """Crawler implementation using Exa API."""

    def __init__(self, config: ExaCrawlerConfig) -> None:
        """Initialize the Exa crawler.

        Args:
            config: Configuration for the crawler
        """
        super().__init__(config)
        self.config: ExaCrawlerConfig = config

    @property
    def needs_parser(self) -> bool:
        return True

    def crawl(self, urls: List[str]) -> List[Document]:
        """Crawl the given URLs using Exa SDK.

        Args:
            urls: List of URLs to crawl

        Returns:
            List of Documents with content extracted from the URLs

        Raises:
            LangroidImportError: If the exa package is not installed
            ValueError: If the Exa API key is not set
        """
        try:
            from exa_py import Exa
        except ImportError:
            raise LangroidImportError("exa", "exa")

        if not self.config.api_key:
            raise ValueError("EXA_API_KEY key is required in your env or .env")

        exa = Exa(self.config.api_key)
        docs = []

        try:
            for url in urls:
                parsed_doc_chunks = self._process_document(url)
                if parsed_doc_chunks:
                    docs.extend(parsed_doc_chunks)
                    continue
                else:
                    results = exa.get_contents(
                        [url],
                        livecrawl="always",
                        text={
                            "include_html_tags": True,
                        },
                    )
                    result = results.results[0]
                    if result.text:
                        md_text = md.markdownify(result.text, heading_style="ATX")
                        # append a NON-chunked document
                        # (metadata.is_chunk = False, so will be chunked downstream)
                        docs.append(
                            Document(
                                content=md_text,
                                metadata=DocMetaData(
                                    source=url,
                                    title=getattr(result, "title", "Unknown Title"),
                                    published_date=getattr(
                                        result, "published_date", "Unknown Date"
                                    ),
                                ),
                            )
                        )

        except Exception as e:
            logging.error(f"Error retrieving content from Exa API: {e}")

        return docs


class Crawl4aiCrawler(BaseCrawler):
    """
    Crawler implementation using the crawl4ai library.

    This crawler intelligently dispatches URLs. Standard web pages are rendered
    and scraped using the crawl4ai browser engine. Direct links to documents
    (PDF, DOCX, etc.) are delegated to the framework's internal DocumentParser.
    """

    def __init__(self, config: Crawl4aiConfig) -> None:
        """Initialize the Crawl4ai crawler."""
        super().__init__(config)
        self.config: Crawl4aiConfig = config

    @property
    def needs_parser(self) -> bool:
        """
        Indicates that this crawler relies on the framework's DocumentParser
        for handling specific file types like PDF, DOCX, etc., which
        the browser engine cannot parse directly.
        """
        return True

    def crawl(self, urls: List[str]) -> List[Document]:
        """
        Executes the crawl by separating document URLs from web page URLs.

        - Document URLs (.pdf, .docx, etc.) are processed using `_process_document`.
        - Web page URLs are handled using the async crawl4ai engine.
        """
        all_documents: List[Document] = []
        webpage_urls: List[str] = []

        # Step 1: Separate URLs into documents and web pages
        for url in urls:
            parsed_doc_chunks = self._process_document(url)
            if parsed_doc_chunks:
                all_documents.extend(parsed_doc_chunks)
            else:
                webpage_urls.append(url)

        # Step 2: Process web page URLs asynchronously
        if webpage_urls:
            try:
                loop = asyncio.get_running_loop()
                if loop.is_running():
                    import nest_asyncio

                    nest_asyncio.apply()
                web_docs = asyncio.run(self._async_crawl(webpage_urls))
            except RuntimeError:
                web_docs = asyncio.run(self._async_crawl(webpage_urls))

            all_documents.extend(web_docs)

        return all_documents

    def _translate_result_to_document(
        self, result: "CrawlResult"
    ) -> Optional[Document]:
        """Converts a crawl4ai CrawlResult into the framework's Document format."""
        if not result.success:
            logging.warning(
                f"Crawl4ai failed for URL {result.url}: {result.error_message}"
            )
            return None

        content = ""
        if result.extracted_content:
            content = result.extracted_content
        elif result.markdown:
            if (
                hasattr(result.markdown, "fit_markdown")
                and result.markdown.fit_markdown
            ):
                content = result.markdown.fit_markdown
            elif hasattr(result.markdown, "raw_markdown"):
                content = result.markdown.raw_markdown
            else:
                content = str(result.markdown)

        if not content:
            logging.warning(f"Crawl4ai returned no content for URL {result.url}")
            return None

        # Extract metadata safely
        title = "Unknown Title"
        published_date = "Unknown Date"

        if result.metadata:
            title = result.metadata.get("title", "Unknown Title")
            # Try common date field names
            for date_field in [
                "published_date",
                "datePublished",
                "article:published_time",
                "pubdate",
            ]:
                if date_field in result.metadata:
                    published_date = result.metadata.get(date_field)
                    break

        meta = DocMetaData(
            source=result.url,
            title=title,
            published_date=published_date,
            # Note: source_content is meant for reference content, not metadata
            # Keeping it minimal as other crawlers don't populate it
        )
        return Document(content=content, metadata=meta)

    async def _async_crawl(self, urls: List[str]) -> List[Document]:
        try:
            from crawl4ai import AsyncWebCrawler

            # Import configs here for lazy loading
            from crawl4ai.async_configs import BrowserConfig, CrawlerRunConfig
        except ImportError:
            raise LangroidImportError(
                "crawl4ai", "pip install 'crawl4ai[all]' or 'crawl4ai'"
            )

        # CHANGE 2: Handle the new optional config fields.
        # Use the user-provided config if it exists, otherwise create a default one.
        browser_config = self.config.browser_config or BrowserConfig()
        run_config = self.config.run_config or CrawlerRunConfig()

        if self.config.extraction_strategy:
            run_config.extraction_strategy = self.config.extraction_strategy
        if self.config.markdown_strategy:
            run_config.markdown_generator = self.config.markdown_strategy
        if self.config.deep_crawl_strategy:
            run_config.deep_crawl_strategy = self.config.deep_crawl_strategy
        if self.config.scraping_strategy:
            run_config.scraping_strategy = self.config.scraping_strategy

        crawled_documents: List[Document] = []

        async with AsyncWebCrawler(config=browser_config) as crawler:
            if self.config.crawl_mode == "simple":
                for url in urls:
                    result = await crawler.arun(url, config=run_config)
                    doc = self._translate_result_to_document(result)
                    if doc:
                        crawled_documents.append(doc)

            elif self.config.crawl_mode == "deep":
                if not urls:
                    return []
                if not run_config.deep_crawl_strategy:
                    logging.warning(
                        "Deep crawl mode requires a deep_crawl_strategy in the config."
                    )
                    return []

                # In deep crawl mode, `crawl4ai` will discover and crawl pages
                # starting from the seed URL. It will not process direct document links
                # found during the deep crawl; it is designed to follow hyperlinks.
                crawl_results = await crawler.arun(urls[0], config=run_config)

                if isinstance(crawl_results, list):
                    for result in crawl_results:
                        doc = self._translate_result_to_document(result)
                        if doc:
                            crawled_documents.append(doc)
                else:
                    async for result in crawl_results:
                        doc = self._translate_result_to_document(result)
                        if doc:
                            crawled_documents.append(doc)

        return crawled_documents


class URLLoader:
    """Loads URLs and extracts text using a specified crawler."""

    def __init__(
        self,
        urls: List[Any],
        parsing_config: ParsingConfig = ParsingConfig(),
        crawler_config: Optional[BaseCrawlerConfig] = None,
    ):
        """Initialize the URL loader.

        Args:
            urls: List of URLs to load
            parsing_config: Configuration for parsing
            crawler_config: Configuration for the crawler
        """
        self.urls = urls
        self.parsing_config = parsing_config

        if crawler_config is None:
            crawler_config = TrafilaturaConfig(parser=Parser(parsing_config))

        self.crawler = CrawlerFactory.create_crawler(crawler_config)
        if self.crawler.needs_parser:
            self.crawler.parser = Parser(parsing_config)

    def load(self) -> List[Document]:
        """Load the URLs using the specified crawler."""
        return self.crawler.crawl(self.urls)
</file>

<file path="langroid/parsing/urls.py">
import logging
import os
import tempfile
import urllib.parse
import urllib.robotparser
from typing import List, Optional, Set, Tuple
from urllib.parse import urldefrag, urljoin, urlparse

import fire
import requests
from bs4 import BeautifulSoup
from pydantic import BaseModel, HttpUrl, TypeAdapter, ValidationError
from rich import print
from rich.prompt import Prompt

logger = logging.getLogger(__name__)


def url_to_tempfile(url: str) -> str:
    """
    Fetch content from the given URL and save it to a temporary local file.

    Args:
        url (str): The URL of the content to fetch.

    Returns:
        str: The path to the temporary file where the content is saved.

    Raises:
        HTTPError: If there's any issue fetching the content.
    """

    response = requests.get(url)
    response.raise_for_status()  # Raise an exception for HTTP errors

    # Create a temporary file and write the content
    with tempfile.NamedTemporaryFile(delete=False, suffix=".tmp") as temp_file:
        temp_file.write(response.content)
        return temp_file.name


def get_user_input(msg: str, color: str = "blue") -> str:
    """
    Prompt the user for input.
    Args:
        msg: printed prompt
        color: color of the prompt
    Returns:
        user input
    """
    color_str = f"[{color}]{msg} " if color else msg + " "
    print(color_str, end="")
    return input("")


def get_list_from_user(
    prompt: str = "Enter input (type 'done' or hit return to finish)",
    n: int | None = None,
) -> List[str]:
    """
    Prompt the user for inputs.
    Args:
        prompt: printed prompt
        n: how many inputs to prompt for. If None, then prompt until done, otherwise
            quit after n inputs.
    Returns:
        list of input strings
    """
    # Create an empty set to store the URLs.
    input_set = set()

    # Use a while loop to continuously ask the user for URLs.
    for _ in range(n or 1000):
        # Prompt the user for input.
        input_str = Prompt.ask(f"[blue]{prompt}")

        # Check if the user wants to exit the loop.
        if input_str.lower() == "done" or input_str == "":
            break

        # if it is a URL, ask how many to crawl
        if is_url(input_str):
            url = input_str
            input_str = Prompt.ask("[blue] How many new URLs to crawl?", default="0")
            max_urls = int(input_str) + 1
            tot_urls = list(find_urls(url, max_links=max_urls, max_depth=2))
            tot_urls_str = "\n".join(tot_urls)
            print(
                f"""
                Found these {len(tot_urls)} links upto depth 2:
                {tot_urls_str}
                """
            )

            input_set.update(tot_urls)
        else:
            input_set.add(input_str.strip())

    return list(input_set)


class Url(BaseModel):
    url: HttpUrl


def is_url(s: str) -> bool:
    try:
        url_adapter = TypeAdapter(HttpUrl)
        Url(url=url_adapter.validate_python(s))
        return True
    except ValidationError:
        return False


def get_urls_paths_bytes_indices(
    inputs: List[str | bytes],
) -> Tuple[List[int], List[int], List[int]]:
    """
    Given a list of inputs, return a
    list of indices of URLs, list of indices of paths, list of indices of byte-contents.
    Args:
        inputs: list of strings or bytes
    Returns:
        list of Indices of URLs,
        list of indices of paths,
        list of indices of byte-contents
    """
    urls = []
    paths = []
    byte_list = []
    for i, item in enumerate(inputs):
        if isinstance(item, bytes):
            byte_list.append(i)
            continue
        try:
            url_adapter = TypeAdapter(HttpUrl)
            Url(url=url_adapter.validate_python(item))
            urls.append(i)
        except ValidationError:
            if os.path.exists(item):
                paths.append(i)
            else:
                logger.warning(f"{item} is neither a URL nor a path.")
    return urls, paths, byte_list


def crawl_url(url: str, max_urls: int = 1) -> List[str]:
    """
    Crawl starting at the url and return a list of URLs to be parsed,
    up to a maximum of `max_urls`.
    This has not been tested to work as intended. Ignore.
    """
    from trafilatura.spider import focused_crawler

    if max_urls == 1:
        # no need to crawl, just return the original list
        return [url]

    to_visit = None
    known_urls = None

    # Create a RobotFileParser object
    robots = urllib.robotparser.RobotFileParser()
    while True:
        if known_urls is not None and len(known_urls) >= max_urls:
            break
        # Set the RobotFileParser object to the website's robots.txt file
        robots.set_url(url + "/robots.txt")
        robots.read()

        if robots.can_fetch("*", url):
            # Start or resume the crawl
            to_visit, known_urls = focused_crawler(
                url,
                max_seen_urls=max_urls,
                max_known_urls=max_urls,
                todo=to_visit,
                known_links=known_urls,
                rules=robots,
            )
        if to_visit is None:
            break

    if known_urls is None:
        return [url]
    final_urls = [s.strip() for s in known_urls]
    return list(final_urls)[:max_urls]


def find_urls(
    url: str = "https://en.wikipedia.org/wiki/Generative_pre-trained_transformer",
    max_links: int = 20,
    visited: Optional[Set[str]] = None,
    depth: int = 0,
    max_depth: int = 2,
    match_domain: bool = True,
) -> Set[str]:
    """
    Recursively find all URLs on a given page.

    Args:
        url (str): The URL to start from.
        max_links (int): The maximum number of links to find.
        visited (set): A set of URLs that have already been visited.
        depth (int): The current depth of the recursion.
        max_depth (int): The maximum depth of the recursion.
        match_domain (bool): Whether to only return URLs that are on the same domain.

    Returns:
        set: A set of URLs found on the page.
    """

    if visited is None:
        visited = set()

    if url in visited or depth > max_depth:
        return visited

    visited.add(url)
    base_domain = urlparse(url).netloc

    try:
        response = requests.get(url, timeout=5)
        response.raise_for_status()
        soup = BeautifulSoup(response.text, "html.parser")
        links = [
            urljoin(url, a["href"])  # type: ignore
            for a in soup.find_all("a", href=True)
        ]

        # Defrag links: discard links that are to portions of same page
        defragged_links = list(
            set(urldefrag(link).url for link in links)  # type: ignore
        )

        # Filter links based on domain matching requirement
        domain_matching_links = [
            link for link in defragged_links if urlparse(link).netloc == base_domain
        ]

        # ensure url is first, since below we are taking first max_links urls
        domain_matching_links = [url] + [x for x in domain_matching_links if x != url]

        # If found links exceed max_links, return immediately
        if len(domain_matching_links) >= max_links:
            return set(domain_matching_links[:max_links])

        for link in domain_matching_links:
            if len(visited) >= max_links:
                break

            if link not in visited:
                visited.update(
                    find_urls(
                        link,
                        max_links,
                        visited,
                        depth + 1,
                        max_depth,
                        match_domain,
                    )
                )

    except (requests.RequestException, Exception) as e:
        print(f"Error fetching {url}. Error: {e}")

    return set(list(visited)[:max_links])


def org_user_from_github(url: str) -> str:
    parsed = urllib.parse.urlparse(url)
    org, user = parsed.path.lstrip("/").split("/")
    return f"{org}-{user}"


if __name__ == "__main__":
    # Example usage
    found_urls = set(fire.Fire(find_urls))
    for url in found_urls:
        print(url)
</file>

<file path="langroid/parsing/utils.py">
import difflib
import logging
import random
import re
from functools import cache
from itertools import islice
from typing import Iterable, List, Sequence, TypeVar

from faker import Faker

from langroid.mytypes import Document
from langroid.parsing.document_parser import DocumentType
from langroid.parsing.parser import Parser, ParsingConfig
from langroid.parsing.repo_loader import RepoLoader
from langroid.parsing.url_loader import URLLoader
from langroid.parsing.urls import get_urls_paths_bytes_indices

Faker.seed(23)
random.seed(43)

logger = logging.getLogger(__name__)


def download_nltk_resource(resource: str) -> None:
    import nltk

    @cache
    def _download() -> None:
        try:
            nltk.data.find(resource)
        except LookupError:
            model = resource.split("/")[-1]
            nltk.download(model, quiet=True)

    _download()


T = TypeVar("T")


def batched(iterable: Iterable[T], n: int) -> Iterable[Sequence[T]]:
    """Batch data into tuples of length n. The last batch may be shorter."""
    # batched('ABCDEFG', 3) --> ABC DEF G
    if n < 1:
        raise ValueError("n must be at least one")
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch


def generate_random_sentences(k: int) -> str:
    # Load the sample text
    import nltk
    from nltk.corpus import gutenberg

    download_nltk_resource("corpora/gutenberg")
    download_nltk_resource("tokenizers/punkt")

    text = gutenberg.raw("austen-emma.txt")

    # Split the text into sentences
    sentences = nltk.tokenize.sent_tokenize(text)

    # Generate k random sentences
    random_sentences = random.choices(sentences, k=k)
    return " ".join(random_sentences)


def generate_random_text(num_sentences: int) -> str:
    fake = Faker()
    text = ""
    for _ in range(num_sentences):
        text += fake.sentence() + " "
    return text


def closest_string(query: str, string_list: List[str]) -> str:
    """Find the closest match to the query in a list of strings.

    This function is case-insensitive and ignores leading and trailing whitespace.
    If no match is found, it returns 'No match found'.

    Args:
        query (str): The string to match.
        string_list (List[str]): The list of strings to search.

    Returns:
        str: The closest match to the query from the list, or 'No match found'
             if no match is found.
    """
    # Create a dictionary where the keys are the standardized strings and
    # the values are the original strings.
    str_dict = {s.lower().strip(): s for s in string_list}

    # Standardize the query and find the closest match in the list of keys.
    closest_match = difflib.get_close_matches(
        query.lower().strip(), str_dict.keys(), n=1
    )

    # Retrieve the original string from the value in the dictionary.
    original_closest_match = (
        str_dict[closest_match[0]] if closest_match else "No match found"
    )

    return original_closest_match


def split_paragraphs(text: str) -> List[str]:
    """
    Split the input text into paragraphs using "\n\n" as the delimiter.

    Args:
        text (str): The input text.

    Returns:
        list: A list of paragraphs.
    """
    # Split based on a newline, followed by spaces/tabs, then another newline.
    paras = re.split(r"\n[ \t]*\n", text)
    return [para.strip() for para in paras if para.strip()]


def split_newlines(text: str) -> List[str]:
    """
    Split the input text into lines using "\n" as the delimiter.

    Args:
        text (str): The input text.

    Returns:
        list: A list of lines.
    """
    lines = re.split(r"\n", text)
    return [line.strip() for line in lines if line.strip()]


def number_segments(s: str, granularity: int = 1) -> str:
    """
    Number the segments in a given text, preserving paragraph structure.
    A segment is a sequence of `len` consecutive "sentences", where a "sentence"
    is either a normal sentence, or if there isn't enough punctuation to properly
    identify sentences, then we use a pseudo-sentence via heuristics (split by newline
    or failing that, just split every 40 words). The goal here is simply to number
    segments at a reasonable granularity so the LLM can identify relevant segments,
    in the RelevanceExtractorAgent.

    Args:
        s (str): The input text.
        granularity (int): The number of sentences in a segment.
            If this is -1, then the entire text is treated as a single segment,
            and is numbered as <#1#>.

    Returns:
        str: The text with segments numbered in the style <#1#>, <#2#> etc.

    Example:
        >>> number_segments("Hello world! How are you? Have a good day.")
        '<#1#> Hello world! <#2#> How are you? <#3#> Have a good day.'
    """
    import nltk

    if granularity < 0:
        return "<#1#> " + s
    numbered_text = []
    count = 0

    paragraphs = split_paragraphs(s)
    for paragraph in paragraphs:
        sentences = nltk.sent_tokenize(paragraph)
        # Some docs are problematic (e.g. resumes) and have no (or too few) periods,
        # so we can't split usefully into sentences.
        # We try a series of heuristics to split into sentences,
        # until the avg num words per sentence is less than 40.
        avg_words_per_sentence = sum(
            len(nltk.word_tokenize(sentence)) for sentence in sentences
        ) / len(sentences)
        if avg_words_per_sentence > 40:
            sentences = split_newlines(paragraph)
        avg_words_per_sentence = sum(
            len(nltk.word_tokenize(sentence)) for sentence in sentences
        ) / len(sentences)
        if avg_words_per_sentence > 40:
            # Still too long, just split on every 40 words
            sentences = []
            for sentence in nltk.sent_tokenize(paragraph):
                words = nltk.word_tokenize(sentence)
                for i in range(0, len(words), 40):
                    # if there are less than 20 words left after this,
                    # just add them to the last sentence and break
                    if len(words) - i < 20:
                        sentences.append(" ".join(words[i:]))
                        break
                    else:
                        sentences.append(" ".join(words[i : i + 40]))
        for i, sentence in enumerate(sentences):
            num = count // granularity + 1
            number_prefix = f"<#{num}#>" if count % granularity == 0 else ""
            sentence = f"{number_prefix} {sentence}"
            count += 1
            sentences[i] = sentence
        numbered_paragraph = " ".join(sentences)
        numbered_text.append(numbered_paragraph)

    return "  \n\n  ".join(numbered_text)


def number_sentences(s: str) -> str:
    return number_segments(s, granularity=1)


def parse_number_range_list(specs: str) -> List[int]:
    """
    Parse a specs string like "3,5,7-10" into a list of integers.

    Args:
        specs (str): A string containing segment numbers and/or ranges
                     (e.g., "3,5,7-10").

    Returns:
        List[int]: List of segment numbers.

    Example:
        >>> parse_number_range_list("3,5,7-10")
        [3, 5, 7, 8, 9, 10]
    """
    spec_indices = set()  # type: ignore
    for part in specs.split(","):
        # some weak LLMs may generate <#1#> instead of 1, so extract just the digits
        # or the "-"
        part = "".join(char for char in part if char.isdigit() or char == "-")
        if "-" in part:
            start, end = map(int, part.split("-"))
            spec_indices.update(range(start, end + 1))
        else:
            spec_indices.add(int(part))

    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.
    """
    paragraphs = split_paragraphs(text)
    cleaned_paragraphs = [" ".join(p.split()) for p in paragraphs if p]
    return "\n\n".join(cleaned_paragraphs)  # Join the cleaned paragraphs.


def extract_numbered_segments(s: str, specs: str) -> str:
    """
    Extract specified segments from a numbered text, preserving paragraph structure.

    Args:
        s (str): The input text containing numbered segments.
        specs (str): A string containing segment numbers and/or ranges
                     (e.g., "3,5,7-10").

    Returns:
        str: Extracted segments, keeping original paragraph structures.

    Example:
        >>> text = "(1) Hello world! (2) How are you? (3) Have a good day."
        >>> extract_numbered_segments(text, "1,3")
        'Hello world! Have a good day.'
    """
    # Use the helper function to get the list of indices from specs
    if specs.strip() == "":
        return ""
    spec_indices = parse_number_range_list(specs)

    # Regular expression to identify numbered segments like
    # <#1#> Hello world! This is me. <#2#> How are you? <#3#> Have a good day.
    # Note we match any character between segment markers, including newlines.
    segment_pattern = re.compile(r"<#(\d+)#>([\s\S]*?)(?=<#\d+#>|$)")

    # Split the text into paragraphs while preserving their boundaries
    paragraphs = split_paragraphs(s)

    extracted_paragraphs = []

    for paragraph in paragraphs:
        segments_with_numbers = segment_pattern.findall(paragraph)

        # Extract the desired segments from this paragraph
        extracted_segments = [
            segment
            for num, segment in segments_with_numbers
            if int(num) in spec_indices
        ]

        # If we extracted any segments from this paragraph,
        # join them with ellipsis (...) and append to results.
        if extracted_segments:
            extracted_paragraphs.append("...".join(extracted_segments))

    return "\n\n".join(extracted_paragraphs)


def extract_content_from_path(
    path: bytes | str | List[bytes | str],
    parsing: ParsingConfig,
    doc_type: str | DocumentType | None = None,
) -> str | List[str]:
    """
    Extract the content from a file path or URL, or a list of file paths or URLs.

    Args:
        path (bytes | str | List[str]): The file path or URL, or a list of file paths or
            URLs, or bytes content. The bytes option is meant to support cases
            where upstream code may have already loaded the content (e.g., from a
            database or API) and we want to avoid having to copy the content to a
            temporary file.
        parsing (ParsingConfig): The parsing configuration.
        doc_type (str | DocumentType | None): The document type if known.
            If multiple paths are given, this MUST apply to ALL docs.

    Returns:
        str | List[str]: The extracted content if a single file path or URL is provided,
                or a list of extracted contents if a
                list of file paths or URLs is provided.
    """
    if isinstance(path, str) or isinstance(path, bytes):
        paths = [path]
    elif isinstance(path, list) and len(path) == 0:
        return ""
    else:
        paths = path

    url_idxs, path_idxs, byte_idxs = get_urls_paths_bytes_indices(paths)
    urls = [paths[i] for i in url_idxs]
    path_list = [paths[i] for i in path_idxs]
    byte_list = [paths[i] for i in byte_idxs]
    path_list.extend(byte_list)
    parser = Parser(parsing)
    docs: List[Document] = []
    try:
        if len(urls) > 0:
            loader = URLLoader(urls=urls, parser=parser)  # type: ignore
            docs = loader.load()
        if len(path_list) > 0:
            for p in path_list:
                path_docs = RepoLoader.get_documents(
                    p, parser=parser, doc_type=doc_type
                )
                docs.extend(path_docs)
    except Exception as e:
        logger.warning(f"Error loading path {paths}: {e}")
        return ""
    if len(docs) == 1:
        return docs[0].content
    else:
        return [d.content for d in docs]
</file>

<file path="langroid/prompts/__init__.py">
from . import dialog
from . import prompts_config
from . import templates

__all__ = [
    "dialog",
    "prompts_config",
    "templates",
]
</file>

<file path="langroid/prompts/dialog.py">
from typing import List


def collate_chat_history(inputs: List[tuple[str, str]]) -> str:
    """
    Collate (human, ai) pairs into a single, string
    Args:
        inputs:
    Returns:
    """
    pairs = [
        f"""Human:{human}
        AI:{ai}
        """
        for human, ai in inputs
    ]
    return "\n".join(pairs)
</file>

<file path="langroid/prompts/prompts_config.py">
from pydantic_settings import BaseSettings


class PromptsConfig(BaseSettings):
    max_tokens: int = 1000  # for output; NOT USED ANYWHERE
</file>

<file path="langroid/prompts/templates.py">
from langroid.utils.constants import NO_ANSWER

EXTRACT_RELEVANT = """
    Here is a passage from a long document, followed by a question. 
    In case the passage contains any text relevant to answer the question, return it 
    verbatim.
    {passage}    
    Question:{question}
    Relevant text, if any: """.strip()

EXTRACTION_PROMPT_GPT4 = f"""
Given the content and question below, extract COMPLETE SENTENCES OR PHRASES 
VERBATIM from the content, that are relevant to answering the question (if such text 
exists), even if it contradicts your knowledge, and even if it is factually incorrect.
Do not  make up an answer that is not supported by the content. 
When you answer, be concise, no need to explain anything. If there is no relevant text,
simply say {NO_ANSWER}.

Content: {{content}}
Question: {{question}}
Relevant text, if any:
"""

EXTRACTION_PROMPT = f"""
    Given the content and question below, extract a COMPLETE SENTENCE verbatim from the 
    content that is relevant to answering the question (if such text exists). Do not 
    make up an answer.
    
    Content: The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in
    Paris, France. It is named after Gustave Eiffel, whose company designed and built
    the tower. It is a recognizable landmark.
    Question: Where is the Eiffel Tower located?
    Relevant text, if any: on the Champ de Mars in Paris, France.
    
    Content: Apples and oranges are both fruits, but differ in taste and texture.
    Apples are sweet and crisp, while oranges are citrusy and juicy. Both are
    nutritious and commonly consumed worldwide.
    Question: What are the similarities between apples and oranges?
    Relevant text, if any: both fruits
    
    Content: The sun rises in the east and sets in the west. It is a source of light
    and warmth for the Earth.
    Question: What is the color of the sun?
    Relevant text, if any: {NO_ANSWER}
    
    Content: {{content}}
    Question: {{question}}
    Relevant text (COMPLETE SENTENCE), if any:
    """.strip()

SUMMARY_ANSWER_PROMPT_GPT4 = f"""

        Use the provided NUMBERED EXTRACTS (with sources)  to answer the QUESTION. 
        If there's not enough information, respond with {NO_ANSWER}. Use ONLY the 
        information in these extracts, even if your answer is factually incorrect.
        and even if the answer contradicts other parts of the document. The only 
        important thing is that your answer is consistent with and supported by the 
        extracts. Compose your complete answer, inserting CITATIONS in MARKDOWN format
        [^i][^j] where i,j,... are the extract NUMBERS you are 
        citing.
        For EXAMPLE your answer might look like this (NOTE HOW multiple citations
        are grouped as [^2][^5]):
        
        <ExampleAnswer>
        Beethoven composed the 9th symphony in 1824.[^1] After that he became deaf
        and could not hear his own music. [^2][^5]. He was a prolific composer and
        wrote many famous pieces.
        </ExampleAnswer>
        
        
        NUMBERED EXTRACTS:
        
        {{extracts}}
        
        QUESTION:
        {{question}}

""".strip()

ANSWER_PROMPT_USE_HISTORY_GPT4 = f"""

        Use ANY of the information earlier, as well as the extracts provided below 
        (with sources)  to answer the question. If there's not 
        enough information, respond with {NO_ANSWER}.
        Use only the information in this conversation or these extracts, 
        even if your answer is factually incorrect, and even 
        if the answer contracts other parts of the document.
        The only important thing is that your answer is 
        consistent with information provided here or earlier.
        Compose your complete answer and cite all supporting sources 
        on a separate separate line as "SOURCE:". 
        When citing a SOURCE: be concise, whether it refers to a source in these 
        extracts, or info provided earlier.
        
        {{extracts}}
        {{question}}
        Answer:   
""".strip()


SUMMARY_ANSWER_PROMPT = f"""
        Use the provided extracts (with sources)  to answer the question. 
        If there's not enough information, respond with {NO_ANSWER}.
        Use only the information in these extracts, even if it contradicts your prior 
        knowledge. Justify your answer by citing your sources, as in these examples:
        
        Extract: The tree species in the garden include oak, maple, and birch.
        Source: https://en.wikipedia.org/wiki/Tree
        Extract: The oak trees are known for their longevity and strength.
        Source: https://en.wikipedia.org/wiki/Oak
        Question: What types of trees are in the garden?
        Answer: The types of trees in the garden include oak, maple, and birch.
        SOURCE: https://en.wikipedia.org/wiki/Tree
        TEXT: The tree species in the garden include oak, maple, and birch.
        
        Extract: The experiment involved three groups: control, low dose, and high 
        dose. 
        Source: https://en.wikipedia.org/wiki/Experiment
        Extract: The high dose group showed significant improvement in symptoms.
        Source: https://en.wikipedia.org/wiki/Experiment
        Extract: The control group did not receive any 
        treatment and served as a baseline.
        Source: https://en.wikipedia.org/wiki/Experiment
        Question: How many groups were involved which group showed significant 
        improvement? 
        Answer: There were three groups and the high dose group showed significant 
        improvement in symptoms.
        SOURCE: https://en.wikipedia.org/wiki/Experiment
        TEXT: The experiment involved three groups: control, low dose, and high dose.
        SOURCE: https://en.wikipedia.org/wiki/Experiment
        TEXT: The high dose group showed significant improvement in symptoms.
        
        
        Extract: The CEO announced several new initiatives during the company meeting.
        Source: https://en.wikipedia.org/wiki/CEO
        Extract: The financial performance of the company has been strong this year.
        Source: https://en.wikipedia.org/wiki/CEO
        Question: What new initiatives did the CEO announce?
        Answer: {NO_ANSWER}
        
        {{extracts}}
        {{question}}
        Answer:
        """.strip()
</file>

<file path="langroid/pydantic_v1/__init__.py">
"""
Compatibility layer for Langroid's Pydantic migration.

IMPORTANT: You are importing from langroid.pydantic_v1 but getting Pydantic v2 classes!
Langroid has fully migrated to Pydantic v2, and this compatibility layer is deprecated.
"""

import warnings
import logging

logger = logging.getLogger(__name__)

# Only show the visual warning, not the standard deprecation warning
# The standard warning is too noisy and shows the import line
logger.warning(
    """
╔════════════════════════════════════════════════════════════════════════╗
║                    ⚠️  DEPRECATION WARNING ⚠️                          ║
╠════════════════════════════════════════════════════════════════════════╣
║                                                                        ║
║  You are importing from langroid.pydantic_v1, but you're actually      ║
║  getting Pydantic v2 classes. Langroid has fully migrated to v2.       ║
║                                                                        ║
║  Please update your imports:                                           ║
║    OLD: from langroid.pydantic_v1 import BaseModel, Field              ║
║    NEW: from pydantic import BaseModel, Field                          ║
║                                                                        ║
║  Also ensure your code uses Pydantic v2 patterns:                      ║
║    • Use model_dump() instead of dict()                                ║
║    • Use model_dump_json() instead of json()                           ║
║    • Use ConfigDict instead of class Config                            ║
║    • Use model_validate() instead of parse_obj()                       ║
║                                                                        ║
║  This compatibility layer will be removed in a future version.         ║
╚════════════════════════════════════════════════════════════════════════╝
"""
)

# Import from pydantic v2 directly (not from pydantic.v1)
# This allows existing code to continue working if it's already v2-compatible
from pydantic import *  # noqa: F403, F401

# BaseSettings has moved in v2, import it explicitly
try:
    from pydantic_settings import BaseSettings  # noqa: F401
except ImportError:
    # Fallback for older pydantic versions
    from pydantic import BaseSettings  # type: ignore[no-redef] # noqa: F401

# Explicitly export all items for mypy
__all__ = [
    "BaseModel",
    "BaseSettings",
    "SettingsConfigDict",
    "Field",
    "ConfigDict",
    "ValidationError",
    "field_validator",
    "model_validator",
    "create_model",
    "HttpUrl",
    "AnyUrl",
    "TypeAdapter",
    "parse_obj_as",
    "validator",
    "root_validator",
]
</file>

<file path="langroid/pydantic_v1/main.py">
"""
Compatibility layer for Pydantic v2 migration.

This module now imports directly from Pydantic v2 since all internal code
has been migrated to use Pydantic v2 patterns.
"""

# Import from pydantic.main but don't trigger the warning again
# The warning is already shown when importing from langroid.pydantic_v1
from pydantic.main import *  # noqa: F403, F401
</file>

<file path="langroid/utils/algorithms/__init__.py">
from . import graph

__all__ = ["graph"]
</file>

<file path="langroid/utils/algorithms/graph.py">
"""
Graph algos.
"""

from typing import Dict, 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


@no_type_check
def components(order: np.ndarray) -> List[List[int]]:
    """
    Find the connected components in an undirected graph represented by a matrix.

    Args:
        order (np.ndarray): A matrix with values 0 or 1 indicating
            undirected graph edges. `order[i][j] = 1` means an edge between `i`
            and `j`, and `0` means no edge.

    Returns:
        List[List[int]]: A list of List where each List contains the indices of
            nodes in the same connected component.

    Example:
        order = np.array([
            [1, 1, 0, 0],
            [1, 1, 1, 0],
            [0, 1, 1, 0],
            [0, 0, 0, 1]
        ])
        components(order)
        # [[0, 1, 2], [3]]
    """

    i2g: Dict[int, int] = {}  # index to group mapping
    next_group = 0
    n = order.shape[0]
    for i in range(n):
        connected_groups = {i2g[j] for j in np.nonzero(order[i, :])[0] if j in i2g}

        # If the node is not part of any group
        # and is not connected to any groups, assign a new group
        if not connected_groups:
            i2g[i] = next_group
            next_group += 1
        else:
            # If the node is connected to multiple groups, we merge them
            main_group = min(connected_groups)
            for j in np.nonzero(order[i, :])[0]:
                if i2g.get(j) in connected_groups:
                    i2g[j] = main_group
            i2g[i] = main_group

    # Convert i2g to a list of Lists
    groups: Dict[int, List[int]] = {}
    for index, group in i2g.items():
        if group not in groups:
            groups[group] = []
        groups[group].append(index)

    return list(groups.values())
</file>

<file path="langroid/utils/output/__init__.py">
from . import printing
from .printing import (
    shorten_text,
    print_long_text,
    show_if_debug,
    SuppressLoggerWarnings,
    PrintColored,
)
from .status import status


__all__ = [
    "printing",
    "shorten_text",
    "print_long_text",
    "show_if_debug",
    "SuppressLoggerWarnings",
    "PrintColored",
    "status",
]
</file>

<file path="langroid/utils/output/citations.py">
import logging
from typing import List, Tuple

from langroid.mytypes import Document

logger = logging.getLogger(__name__)


def extract_markdown_references(md_string: str) -> List[int]:
    """
    Extracts markdown references (e.g., [^1], [^2]) from a string and returns
    them as a sorted list of integers.

    Args:
        md_string (str): The markdown string containing references.

    Returns:
        list[int]: A sorted list of unique integers from the markdown references.
    """
    import re

    # Regex to find all occurrences of [^<number>]
    matches = re.findall(r"\[\^(\d+)\]", md_string)
    # Convert matches to integers, remove duplicates with set, and sort
    return sorted(set(int(match) for match in matches))


def invalid_markdown_citations(md_string: str) -> List[str]:
    """
    Finds non-numeric markdown citations (e.g., [^a], [^xyz]) in a string.

    Args:
        md_string (str): The markdown string to search for invalid citations.

    Returns:
        List[str]: List of invalid citation strings (without brackets/caret).
    """
    import re

    # Find all citation references first
    matches = re.findall(r"\[\^([^\]\s]+)\]", md_string)

    # Filter out purely numeric citations
    invalid_citations = [match for match in matches if not match.isdigit()]

    return sorted(set(invalid_citations))


def format_footnote_text(content: str, width: int = 0) -> str:
    """
    Formats the content so that each original line is individually processed.
    - If width=0, no wrapping is done (lines remain as is).
    - If width>0, lines are wrapped to that width.
    - Blank lines remain blank (with indentation).
    - Everything is indented by 4 spaces (for markdown footnotes).

    Args:
        content (str): The text of the footnote to be formatted.
        width (int): Maximum width of the text lines. If 0, lines are not wrapped.

    Returns:
        str: Properly formatted markdown footnote text.
    """
    import textwrap

    indent = "    "  # 4 spaces for markdown footnotes
    lines = content.split("\n")  # keep original line structure

    output_lines = []
    for line in lines:
        # If the line is empty (or just spaces), keep it blank (but indented)
        if not line.strip():
            output_lines.append(indent)
            continue

        if width > 0:
            # Wrap each non-empty line to the specified width
            wrapped = textwrap.wrap(line, width=width)
            if not wrapped:
                # If textwrap gives nothing, add a blank (indented) line
                output_lines.append(indent)
            else:
                for subline in wrapped:
                    output_lines.append(indent + subline)
        else:
            # No wrapping: just indent the original line
            output_lines.append(indent + line)

    # Join them with newline so we preserve the paragraph/blank line structure
    return "\n".join(output_lines)


def format_cited_references(
    citations: List[int], passages: list[Document]
) -> Tuple[str, str]:
    """
    Given a list of (integer) citations, and a list of passages, return a string
    that can be added as a footer to the main text, to show sources cited.

    Args:
        citations (list[int]): list of citations, presumably from main text
        passages (list[Document]): list of passages (Document objects)

    Returns:
        str: formatted string of FULL citations (i.e. reference AND content)
            for footnote in markdown;
        str: formatted string of BRIEF citations (i.e. reference only)
            for footnote in markdown.
    """
    citations_str = ""
    full_citations_str = ""
    if len(citations) > 0:
        # append [i] source, content for each citation
        good_citations = [c for c in citations if c > 0 and c <= len(passages)]
        if len(good_citations) < len(citations):
            logger.warning(f"Invalid citations: {set(citations) - set(good_citations)}")

        # source and content for each citation
        full_citations_str = "\n".join(
            [
                f"[^{c}] {str(passages[c-1].metadata)}"
                f"\n{format_footnote_text(passages[c-1].content)}"
                for c in good_citations
            ]
        )

        # source for each citation
        citations_str = "\n".join(
            [f"[^{c}] {str(passages[c-1].metadata)}" for c in good_citations]
        )
    return full_citations_str, citations_str
</file>

<file path="langroid/utils/output/printing.py">
import logging
import sys
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Type

from rich import print as rprint
from rich.text import Text

from langroid.utils.configuration import settings
from langroid.utils.constants import Colors


def shorten_text(text: str, chars: int = 40) -> str:
    text = " ".join(text.split())
    return text[:chars] + "..." + text[-chars:] if len(text) > 2 * chars else text


def print_long_text(
    color: str, style: str, preamble: str, text: str, chars: Optional[int] = None
) -> None:
    if chars is not None:
        text = " ".join(text.split())
        text = text[:chars] + "..." + text[-chars:] if len(text) > 2 * chars else text
    styled_text = Text(text, style=style)
    rprint(f"[{color}]{preamble} {styled_text}")


def show_if_debug(
    text: str,
    preamble: str,
    chars: Optional[int] = None,
    color: str = "red",
    style: str = "italic",
) -> None:
    if settings.debug:
        print_long_text(color, style, preamble, text, chars)


class PrintColored:
    """Context to temporarily print in a desired color"""

    def __init__(self, color: str):
        self.color = color

    def __enter__(self) -> None:
        sys.stdout.write(self.color)
        sys.stdout.flush()

    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()


class SuppressLoggerWarnings:
    def __init__(self, logger: str | None = None):
        # If no logger name is given, get the root logger
        self.logger = logging.getLogger(logger)
        self.original_level = self.logger.getEffectiveLevel()

    def __enter__(self) -> None:
        # Set the logging level to 'ERROR' to suppress warnings
        self.logger.setLevel(logging.ERROR)

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Any,
    ) -> None:
        # Reset the logging level to its original value
        self.logger.setLevel(self.original_level)
</file>

<file path="langroid/utils/output/status.py">
import logging
from contextlib import AbstractContextManager, ExitStack
from typing import Any

from rich.console import Console
from rich.errors import LiveError

from langroid.utils.configuration import quiet_mode, settings

console = Console()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


def status(
    msg: str,
    log_if_quiet: bool = True,
) -> AbstractContextManager[Any]:
    """
    Displays a rich spinner if not in quiet mode, else optionally logs the message.
    """
    stack = ExitStack()
    logged = False
    if settings.quiet and log_if_quiet:
        logged = True
        logger.info(msg)

    if not settings.quiet:
        try:
            stack.enter_context(console.status(msg))
        except LiveError:
            if not logged:
                logger.info(msg)

    # When using rich spinner, we enforce quiet mode
    # (since output will be messy otherwise);
    # We make an exception to this when debug is enabled.
    stack.enter_context(quiet_mode(not settings.debug))

    return stack
</file>

<file path="langroid/utils/__init__.py">
from . import configuration
from . import globals
from . import constants
from . import logging
from . import pydantic_utils
from . import system
from . import output
from . import object_registry

__all__ = [
    "configuration",
    "globals",
    "constants",
    "logging",
    "pydantic_utils",
    "system",
    "output",
    "object_registry",
]
</file>

<file path="langroid/utils/configuration.py">
import os
import threading
from contextlib import contextmanager
from typing import Any, Dict, Iterator, List, Literal, cast

from dotenv import find_dotenv, load_dotenv
from pydantic_settings import BaseSettings, SettingsConfigDict

# Global reentrant lock to serialize any modifications to the global settings.
_global_lock = threading.RLock()


class Settings(BaseSettings):
    debug: bool = False  # show debug messages?
    max_turns: int = -1  # maximum number of turns in a task (to avoid inf loop)
    progress: bool = False  # show progress spinners/bars?
    stream: bool = True  # stream output?
    cache: bool = True  # use cache?
    cache_type: Literal["redis", "fakeredis", "none"] = "redis"  # cache type
    chat_model: str = ""  # language model name, e.g. litellm/ollama/llama2
    quiet: bool = False  # quiet mode (i.e. suppress all output)?
    notebook: bool = False  # running in a notebook?

    model_config = SettingsConfigDict(extra="forbid")


# Load environment variables from .env file.
load_dotenv(find_dotenv(usecwd=True))

# The global (default) settings instance.
# This is updated by update_global_settings() and set_global().
_global_settings = Settings()

# Thread-local storage for temporary (per-thread) settings overrides.
_thread_local = threading.local()


class SettingsProxy:
    """
    A proxy for the settings that returns a thread‐local override if set,
    or else falls back to the global settings.
    """

    def __getattr__(self, name: str) -> Any:
        # If the calling thread has set an override, use that.
        if hasattr(_thread_local, "override"):
            return getattr(_thread_local.override, name)
        return getattr(_global_settings, name)

    def __setattr__(self, name: str, value: Any) -> None:
        # All writes go to the global settings.
        setattr(_global_settings, name, value)

    def update(self, new_settings: Settings) -> None:
        _global_settings.__dict__.update(new_settings.__dict__)

    def dict(self) -> Dict[str, Any]:
        # Return a dict view of the settings as seen by the caller.
        # Note that temporary overrides are not “merged” with global settings.
        if hasattr(_thread_local, "override"):
            return cast(
                Dict[str, Any], cast(Settings, _thread_local.override.model_dump())
            )
        return _global_settings.model_dump()


settings = SettingsProxy()


def update_global_settings(cfg: BaseSettings, keys: List[str]) -> None:
    """
    Update global settings so that modules can later access them via, e.g.,

        from langroid.utils.configuration import settings
        if settings.debug: ...

    This updates the global default.
    """
    config_dict = cfg.model_dump()
    filtered_config = {key: config_dict[key] for key in keys if key in config_dict}
    new_settings = Settings(**filtered_config)
    _global_settings.__dict__.update(new_settings.__dict__)


def set_global(key_vals: Settings) -> None:
    """
    Update the global settings object.
    """
    _global_settings.__dict__.update(key_vals.__dict__)


@contextmanager
def temporary_settings(temp_settings: Settings) -> Iterator[None]:
    """
    Temporarily override the settings for the calling thread.

    Within the context, any access to "settings" will use the provided temporary
    settings. Once the context is exited, the thread reverts to the global settings.
    """
    saved = getattr(_thread_local, "override", None)
    _thread_local.override = temp_settings
    try:
        yield
    finally:
        if saved is not None:
            _thread_local.override = saved
        else:
            del _thread_local.override


@contextmanager
def quiet_mode(quiet: bool = True) -> Iterator[None]:
    """
    Temporarily override settings.quiet for the current thread.
    This implementation builds on the thread‑local temporary_settings context manager.
    The effective quiet mode is merged:
    if quiet is already True (from an outer context),
    then it remains True even if a nested context passes quiet=False.
    """
    current_effective = (
        settings.model_dump()
    )  # get the current thread's effective settings
    # Create a new settings instance from the current effective state.
    temp = Settings(**current_effective)
    # Merge the new flag: once quiet is enabled, it stays enabled.
    temp.quiet = settings.quiet or quiet
    with temporary_settings(temp):
        yield


def set_env(settings_instance: BaseSettings) -> None:
    """
    Set environment variables from a BaseSettings instance.

    Each field in the settings is written to os.environ.
    """
    for field_name, field in settings_instance.__class__.model_fields.items():
        env_var_name = field.alias or field_name.upper()
        os.environ[env_var_name] = str(settings_instance.model_dump()[field_name])
</file>

<file path="langroid/utils/constants.py">
from pydantic import BaseModel


# Define the ANSI escape sequences for various colors and reset
class Colors(BaseModel):
    RED: str = "\033[31m"
    BLUE: str = "\033[34m"
    GREEN: str = "\033[32m"
    GREEN_DIMMER: str = "\033[38;5;22m"  # very dark green
    GREEN_DIM: str = "\033[38;5;28m"  # medium-dim green
    ORANGE: str = "\033[33m"  # no standard ANSI color for orange; using yellow
    CYAN: str = "\033[36m"
    MAGENTA: str = "\033[35m"
    YELLOW: str = "\033[33m"
    RESET: str = "\033[0m"


NO_ANSWER = "DO-NOT-KNOW"
DONE = "DONE"
USER_QUIT_STRINGS = ["q", "x", "quit", "exit", "bye", DONE]
PASS = "__PASS__"
PASS_TO = PASS + ":"
SEND_TO = "__SEND__:"
TOOL = "TOOL"
# This is a recommended setting for TaskConfig.addressing_prefix if using it at all;
# prefer to use `RecipientTool` to allow agents addressing others.
# Caution the AT string should NOT contain any 'word' characters, i.e.
# it no letters, digits or underscores.
# See tests/main/test_msg_routing for example usage
AT = "|@|"
TOOL_BEGIN = "TOOL_BEGIN"
TOOL_END = "TOOL_END"
</file>

<file path="langroid/utils/git_utils.py">
import fnmatch
import logging
import textwrap
from pathlib import Path
from typing import List

import git
from github import Github, GithubException

from langroid.utils.system import create_file

logger = logging.getLogger(__name__)


def git_read_file(repo: str, filepath: str) -> str:
    """
    Read the contents of a file from a GitHub repository.

    Args:
        repo (str): The GitHub repository in the format "owner/repo"
        filepath (str): The file path relative to the repository root

    Returns:
        str: The contents of the file as a string
    """
    try:
        g = Github()
        github_repo = g.get_repo(repo)
        file_content = github_repo.get_contents(filepath)
        if isinstance(file_content, list) and len(file_content) > 0:
            return file_content[0].decoded_content.decode("utf-8")
        elif hasattr(file_content, "decoded_content"):
            return file_content.decoded_content.decode("utf-8")
        else:
            logger.error(f"Unexpected file_content type: {type(file_content)}")
            return ""
    except GithubException as e:
        logger.error(f"An error occurred while reading file {filepath}: {e}")
        return ""


def get_file_list(repo: str, dir: str, pat: str = "") -> List[str]:
    """
    Get a list of files in a specified directory of a GitHub repository.

    Args:
        repo (str): The GitHub repository in the format "owner/repo"
        dir (str): The directory path relative to the repository root
        pat (str): Optional wildcard pattern to filter file names (default: "")

    Returns:
        List[str]: A list of file paths in the specified directory
    """
    try:
        g = Github()
        github_repo = g.get_repo(repo)
        contents = github_repo.get_contents(dir)

        file_list = []
        if isinstance(contents, list):
            file_list = [content.path for content in contents if content.type == "file"]
        elif hasattr(contents, "path") and hasattr(contents, "type"):
            if contents.type == "file":
                file_list = [contents.path]

        if pat:
            file_list = [file for file in file_list if fnmatch.fnmatch(file, pat)]
        return sorted(file_list)

    except GithubException as e:
        logger.error(f"An error occurred while fetching file list: {e}")
        return []


def git_init_repo(dir: str) -> git.Repo | None:
    """
    Set up a Git repository in the specified directory.

    Args:
        dir (str): Path to the directory where the Git repository should be initialized

    Returns:
        git.Repo: The initialized Git repository object
    """
    repo_path = Path(dir).expanduser()
    try:
        repo = git.Repo.init(repo_path)
        logger.info(f"Git repository initialized in {repo_path}")

        gitignore_content = textwrap.dedent(
            """
        /target/
        **/*.rs.bk
        Cargo.lock
        """
        ).strip()

        gitignore_path = repo_path / ".gitignore"
        create_file(gitignore_path, gitignore_content)
        logger.info(f"Created .gitignore file in {repo_path}")

        # Ensure the default branch is 'main'
        # Check if we're on the master branch
        if repo.active_branch.name == "master":
            # Rename the branch
            repo.git.branch("-m", "master", "main")
            print("Branch renamed from 'master' to 'main'")
        else:
            print("Current branch is not 'master'. No changes made.")
        return repo
    except git.GitCommandError as e:
        logger.error(f"An error occurred while initializing the repository: {e}")
        return None


def git_commit_file(repo: git.Repo, filepath: str, msg: str) -> None:
    """
    Commit a file to a Git repository.

    Args:
        repo (git.Repo): The Git repository object
        filepath (str): Path to the file to be committed
        msg (str): The commit message

    Returns:
        None
    """
    try:
        repo.index.add([filepath])
        commit_msg = msg or f"Updated {filepath}"
        repo.index.commit(commit_msg)
        logger.info(f"Successfully committed {filepath}: {commit_msg}")
    except git.GitCommandError as e:
        logger.error(f"An error occurred while committing: {e}")


def git_commit_mods(repo: git.Repo, msg: str = "commit all changes") -> None:
    """
    Commit all modifications in the Git repository.
    Does not raise an error if there's nothing to commit.

    Args:
        repo (git.Repo): The Git repository object

    Returns:
        None
    """
    try:
        if repo.is_dirty():
            repo.git.add(update=True)
            repo.index.commit(msg)
            logger.info("Successfully committed all modifications")
        else:
            logger.info("No changes to commit")
    except git.GitCommandError as e:
        logger.error(f"An error occurred while committing modifications: {e}")


def git_restore_repo(repo: git.Repo) -> None:
    """
    Restore all unstaged, uncommitted changes in the Git repository.
    This function undoes any dirty files to the last commit.

    Args:
        repo (git.Repo): The Git repository object

    Returns:
        None
    """
    try:
        if repo.is_dirty():
            repo.git.restore(".")
            logger.info("Successfully restored all unstaged changes")
        else:
            logger.info("No unstaged changes to restore")
    except git.GitCommandError as e:
        logger.error(f"An error occurred while restoring changes: {e}")


def git_restore_file(repo: git.Repo, file_path: str) -> None:
    """
    Restore a specific file in the Git repository to its state in the last commit.
    This function undoes changes to the specified file.

    Args:
        repo (git.Repo): The Git repository object
        file_path (str): Path to the file to be restored

    Returns:
        None
    """
    try:
        repo.git.restore(file_path)
        logger.info(f"Successfully restored file: {file_path}")
    except git.GitCommandError as e:
        logger.error(f"An error occurred while restoring file {file_path}: {e}")


def git_create_checkout_branch(repo: git.Repo, branch: str) -> None:
    """
    Create and checkout a new branch in the given Git repository.
    If the branch already exists, it will be checked out.
    If we're already on the specified branch, no action is taken.

    Args:
        repo (git.Repo): The Git repository object
        branch (str): The name of the branch to create or checkout

    Returns:
        None
    """
    try:
        if repo.active_branch.name == branch:
            logger.info(f"Already on branch: {branch}")
            return

        if branch in repo.heads:
            repo.heads[branch].checkout()
            logger.info(f"Checked out existing branch: {branch}")
        else:
            new_branch = repo.create_head(branch)
            new_branch.checkout()
            logger.info(f"Created and checked out new branch: {branch}")
    except git.GitCommandError as e:
        logger.error(f"An error occurred while creating/checking out branch: {e}")


def git_diff_file(repo: git.Repo, filepath: str) -> str:
    """
    Show diffs of file between the latest commit and the previous one if any.

    Args:
        repo (git.Repo): The Git repository object
        filepath (str): Path to the file to be diffed

    Returns:
        str: The diff output as a string
    """
    try:
        # Get the two most recent commits
        commits = list(repo.iter_commits(paths=filepath, max_count=2))

        if len(commits) < 2:
            return "No previous commit found for comparison."

        # Get the diff between the two commits for the specific file
        diff = repo.git.diff(commits[1].hexsha, commits[0].hexsha, filepath)

        return str(diff)
    except git.GitCommandError as e:
        logger.error(f"An error occurred while getting diff: {e}")
        return f"Error: {str(e)}"
</file>

<file path="langroid/utils/globals.py">
from typing import Any, Dict, Optional, Type, TypeVar, cast

from pydantic import BaseModel
from pydantic.fields import ModelPrivateAttr
from pydantic_core import PydanticUndefined

T = TypeVar("T", bound="GlobalState")


class GlobalState(BaseModel):
    """A base Pydantic model for global states."""

    _instance: Optional["GlobalState"] = None

    @classmethod
    def get_instance(cls: Type["GlobalState"]) -> "GlobalState":
        """
        Get the global instance of the specific subclass.

        Returns:
            The global instance of the subclass.
        """
        # Get the actual value from ModelPrivateAttr when accessing on class
        instance_attr = getattr(cls, "_instance", None)
        actual_instance: Optional["GlobalState"]
        if isinstance(instance_attr, ModelPrivateAttr):
            default_value = instance_attr.default
            if default_value is PydanticUndefined:
                actual_instance = None
            else:
                actual_instance = cast(Optional["GlobalState"], default_value)
        else:
            actual_instance = instance_attr

        if actual_instance is None:
            new_instance = cls()
            cls._instance = new_instance
            return new_instance
        return actual_instance  # type: ignore

    @classmethod
    def set_values(cls: Type[T], **kwargs: Dict[str, Any]) -> None:
        """
        Set values on the global instance of the specific subclass.

        Args:
            **kwargs: The fields and their values to set.
        """
        instance = cls.get_instance()
        for key, value in kwargs.items():
            setattr(instance, key, value)

    @classmethod
    def get_value(cls: Type[T], name: str) -> Any:
        """
        Retrieve the value of a specific field from the global instance.

        Args:
            name (str): The name of the field to retrieve.

        Returns:
            str: The value of the specified field.
        """
        instance = cls.get_instance()
        return getattr(instance, name)
</file>

<file path="langroid/utils/html_logger.py">
"""HTML Logger for Langroid Task System.

This module provides an HTML logger that creates self-contained HTML files
with collapsible log entries for better visualization of agent interactions.
"""

import html
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

from pydantic import BaseModel

from langroid.utils.logging import setup_logger


class HTMLLogger:
    """Logger that outputs task logs as interactive HTML files."""

    def __init__(
        self,
        filename: str,
        log_dir: str = "logs",
        model_info: str = "",
        append: bool = False,
    ):
        """Initialize the HTML logger.

        Args:
            filename: Base name for the log file (without extension)
            log_dir: Directory to store log files
            model_info: Information about the model being used
            append: Whether to append to existing file
        """
        self.filename = filename
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.file_path = self.log_dir / f"{filename}.html"
        self.model_info = model_info
        self.entries: List[Dict[str, Any]] = []
        self.entry_counter = 0
        self.tool_counter = 0

        # Logger for errors
        self.logger = setup_logger(__name__)

        if not append or not self.file_path.exists():
            self._write_header()

    def _write_header(self) -> None:
        """Write the HTML header with CSS and JavaScript."""
        timestamp = datetime.now().strftime("%m/%d/%Y, %I:%M:%S %p")

        html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <meta http-equiv="refresh" content="2">
    <title>{self.filename} - Langroid Task Log</title>
    <style>
        body {{
            background-color: #1e1e1e;
            color: #f0f0f0;
            font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
            font-size: 14px;
            margin: 0;
            padding: 20px;
            line-height: 1.6;
        }}
        
        .header {{
            border: 2px solid #d4a017;
            padding: 15px;
            margin-bottom: 20px;
            color: #d4a017;
            background-color: #2b2b2b;
            border-radius: 5px;
        }}
        
        .header-line {{
            display: flex;
            justify-content: space-between;
            align-items: center;
        }}
        
        .separator {{
            border-bottom: 2px solid #d4a017;
            margin: 20px 0;
        }}
        
        .controls {{
            margin-bottom: 20px;
        }}
        
        .controls {{
            display: flex;
            align-items: center;
            gap: 20px;
        }}
        
        .controls button {{
            background-color: #333;
            color: #f0f0f0;
            border: 1px solid #555;
            padding: 8px 16px;
            cursor: pointer;
            border-radius: 3px;
            font-family: inherit;
        }}
        
        .controls button:hover {{
            background-color: #444;
            border-color: #d4a017;
        }}
        
        .controls label {{
            color: #f0f0f0;
            display: flex;
            align-items: center;
            gap: 8px;
            cursor: pointer;
        }}
        
        .controls input[type="checkbox"] {{
            cursor: pointer;
        }}
        
        .hidden {{
            display: none !important;
        }}
        
        .entry {{
            margin-bottom: 15px;
            padding-left: 10px;
        }}
        
        .entry.faded {{
            opacity: 0.4;
        }}
        
        .entry.important {{
            opacity: 1.0;
        }}
        
        .entry.user .entity-header {{
            color: #00bfff;
        }}
        
        .entry.assistant .entity-header {{
            color: #ff6b6b;
        }}
        
        .entry.llm .entity-header {{
            color: #00ff00;
        }}
        
        .entry.agent .entity-header {{
            color: #ff9500;
        }}
        
        .entry.system .entity-header {{
            color: #888;
        }}
        
        .entry.other .entity-header {{
            color: #999;
        }}
        
        .entity-header {{
            font-weight: bold;
            margin-bottom: 5px;
            cursor: pointer;
        }}
        
        .entity-header:hover {{
            opacity: 0.8;
        }}
        
        .header-main {{
            /* Removed text-transform to preserve tool name casing */
            display: inline;
        }}
        
        .header-content {{
            margin-left: 30px;
            opacity: 0.7;
            font-weight: normal;
            font-style: italic;
            display: block;
        }}
        
        .entry-content {{
            margin-left: 20px;
            margin-top: 5px;
        }}
        
        .entry-content.collapsed {{
            display: none;
        }}
        
        .collapsible {{
            margin: 5px 0;
            margin-left: 20px;
        }}
        
        .toggle {{
            cursor: pointer;
            user-select: none;
            color: #00ff00;
            display: inline-block;
            width: 25px;
            font-family: monospace;
            margin-right: 5px;
        }}
        
        .toggle:hover {{
            color: #00ff00;
            text-shadow: 0 0 5px #00ff00;
        }}
        
        .content {{
            margin-left: 25px;
            margin-top: 5px;
            white-space: pre-wrap;
            word-wrap: break-word;
        }}
        
        .main-content {{
            margin-top: 10px;
            white-space: pre-wrap;
            word-wrap: break-word;
        }}
        
        .collapsed .content {{
            display: none;
        }}
        
        .tool-section {{
            margin: 10px 0;
            margin-left: 20px;
        }}
        
        .tool-name {{
            color: #d4a017;
            font-weight: bold;
        }}
        
        .tool-result {{
            margin-left: 25px;
        }}
        
        .tool-result.success {{
            color: #00ff00;
        }}
        
        .tool-result.error {{
            color: #ff0000;
        }}
        
        .code-block {{
            background-color: #2b2b2b;
            border: 1px solid #444;
            padding: 10px;
            margin: 5px 0;
            border-radius: 3px;
            overflow-x: auto;
        }}
        
        .metadata {{
            color: #888;
            font-size: 0.9em;
            margin-left: 25px;
        }}
        
        
        pre {{
            margin: 0;
            white-space: pre-wrap;
            word-wrap: break-word;
        }}
    </style>
    <script>
        function toggleEntry(entryId) {{
            const contentElement = document.getElementById(entryId + '_content');
            const toggleElement = document.querySelector(
                '#' + entryId + ' .entity-header .toggle'
            );
            
            if (!contentElement || !toggleElement) return;
            
            if (contentElement.classList.contains('collapsed')) {{
                contentElement.classList.remove('collapsed');
                toggleElement.textContent = '[-]';
                // Save expanded state
                localStorage.setItem('expanded_' + entryId, 'true');
            }} else {{
                contentElement.classList.add('collapsed');
                toggleElement.textContent = '[+]';
                // Save collapsed state
                localStorage.setItem('expanded_' + entryId, 'false');
            }}
        }}
        
        function toggle(id) {{
            const element = document.getElementById(id);
            if (!element) return;
            
            element.classList.toggle('collapsed');
            const toggle = element.querySelector('.toggle');
            if (toggle) {{
                toggle.textContent = element.classList.contains('collapsed') 
                    ? '[+]' : '[-]';
            }}
            
            // Save collapsed state for collapsible sections
            localStorage.setItem(
                'collapsed_' + id, element.classList.contains('collapsed')
            );
        }}
        
        let allExpanded = false;
        
        function toggleAll() {{
            const btn = document.getElementById('toggleAllBtn');
            if (allExpanded) {{
                collapseAll();
                btn.textContent = 'Expand All';
                allExpanded = false;
            }} else {{
                expandAll();
                btn.textContent = 'Collapse All';
                allExpanded = true;
            }}
        }}
        
        function expandAll() {{
            // Expand all visible main entries
            const entries = document.querySelectorAll(
                '.entry:not(.hidden) .entry-content'
            );
            entries.forEach(element => {{
                element.classList.remove('collapsed');
            }});
            
            // Update all visible main entry toggles
            const entryToggles = document.querySelectorAll(
                '.entry:not(.hidden) .entity-header .toggle'
            );
            entryToggles.forEach(toggle => {{
                toggle.textContent = '[-]';
            }});
            
            // Expand all visible sub-sections
            const collapsibles = document.querySelectorAll(
                '.entry:not(.hidden) .collapsible'
            );
            collapsibles.forEach(element => {{
                element.classList.remove('collapsed');
                const toggle = element.querySelector('.toggle');
                if (toggle) {{
                    toggle.textContent = '[-]';
                }}
            }});
        }}
        
        function collapseAll() {{
            // Collapse all visible entries
            const entries = document.querySelectorAll(
                '.entry:not(.hidden) .entry-content'
            );
            entries.forEach(element => {{
                element.classList.add('collapsed');
            }});
            
            // Update all visible entry toggles
            const entryToggles = document.querySelectorAll(
                '.entry:not(.hidden) .entity-header .toggle'
            );
            entryToggles.forEach(toggle => {{
                toggle.textContent = '[+]';
            }});
            
            // Collapse all visible sub-sections
            const collapsibles = document.querySelectorAll(
                '.entry:not(.hidden) .collapsible'
            );
            collapsibles.forEach(element => {{
                element.classList.add('collapsed');
                const toggle = element.querySelector('.toggle');
                if (toggle) {{
                    toggle.textContent = '[+]';
                }}
            }});
        }}
        
        function filterEntries() {{
            const checkbox = document.getElementById('filterCheckbox');
            const entries = document.querySelectorAll('.entry');
            
            // Save checkbox state to localStorage
            localStorage.setItem('filterImportant', checkbox.checked);
            
            if (checkbox.checked) {{
                // Show only important entries
                entries.forEach(entry => {{
                    const isImportant = entry.classList.contains('important');
                    if (isImportant) {{
                        entry.classList.remove('hidden');
                    }} else {{
                        entry.classList.add('hidden');
                    }}
                }});
            }} else {{
                // Show all entries
                entries.forEach(entry => {{
                    entry.classList.remove('hidden');
                }});
            }}
            
            // Reset toggle button state
            allExpanded = false;
            document.getElementById('toggleAllBtn').textContent = 'Expand All';
        }}
        
        // Initialize all as collapsed on load
        document.addEventListener('DOMContentLoaded', function() {{
            collapseAll();
            
            // Restore checkbox state from localStorage
            const checkbox = document.getElementById('filterCheckbox');
            const savedState = localStorage.getItem('filterImportant');
            if (savedState !== null) {{
                // Use saved state if it exists
                checkbox.checked = savedState === 'true';
            }}
            // Apply filter based on checkbox state (default is checked)
            if (checkbox.checked) {{
                filterEntries();
            }}
            
            // Restore expanded states from localStorage
            const entries = document.querySelectorAll('.entry');
            entries.forEach(entry => {{
                const entryId = entry.id;
                const expandedState = localStorage.getItem('expanded_' + entryId);
                if (expandedState === 'true') {{
                    const contentElement = document.getElementById(
                        entryId + '_content'
                    );
                    const toggleElement = entry.querySelector('.entity-header .toggle');
                    if (contentElement && toggleElement) {{
                        contentElement.classList.remove('collapsed');
                        toggleElement.textContent = '[-]';
                    }}
                }}
            }});
            
            // Restore collapsible section states
            const collapsibles = document.querySelectorAll('.collapsible');
            collapsibles.forEach(collapsible => {{
                const id = collapsible.id;
                const collapsedState = localStorage.getItem('collapsed_' + id);
                if (collapsedState === 'false') {{
                    collapsible.classList.remove('collapsed');
                    const toggle = collapsible.querySelector('.toggle');
                    if (toggle) {{
                        toggle.textContent = '[-]';
                    }}
                }}
            }});
        }});
    </script>
</head>
<body>
    <div class="header">
        <div class="header-line">
            <div>{self.filename}</div>
            <div id="timestamp">{timestamp}</div>
        </div>
    </div>
    
    <div class="separator"></div>
    
    <div class="controls">
        <button id="toggleAllBtn" onclick="toggleAll()">Expand All</button>
        <label style="margin-left: 20px;">
            <input type="checkbox" id="filterCheckbox" 
                   onchange="filterEntries()" checked>
            Show only important responses
        </label>
    </div>
    
    <div id="content">
"""
        try:
            with open(self.file_path, "w", encoding="utf-8") as f:
                f.write(html_content)
        except Exception as e:
            self.logger.error(f"Failed to write HTML header: {e}")

    def log(self, fields: BaseModel) -> None:
        """Log a message entry.

        Args:
            fields: ChatDocLoggerFields containing all log information
        """
        try:
            entry_html = self._format_entry(fields)
            self._append_to_file(entry_html)
            self.entry_counter += 1
        except Exception as e:
            self.logger.error(f"Failed to log entry: {e}")

    def _format_entry(self, fields: BaseModel) -> str:
        """Format a log entry as HTML.

        Args:
            fields: ChatDocLoggerFields containing all log information

        Returns:
            HTML string for the entry
        """
        entry_id = f"entry_{self.entry_counter}"

        # Get all relevant fields
        responder = str(getattr(fields, "responder", "UNKNOWN"))
        task_name = getattr(fields, "task_name", "root")
        # TODO (CLAUDE) display sender_entity in parens right after responder,
        # other than LLM, e.g. AGENT (USER)
        sender_entity = str(getattr(fields, "sender_entity", ""))
        tool = getattr(fields, "tool", "")
        tool_type = getattr(fields, "tool_type", "")
        content = getattr(fields, "content", "")
        recipient = getattr(fields, "recipient", "")

        # Determine CSS class based on responder
        responder_upper = responder.upper()
        if "USER" in responder_upper:
            css_class = "user"
        elif "LLM" in responder_upper:
            css_class = "llm"
        elif "AGENT" in responder_upper:
            css_class = "agent"
        elif "SYSTEM" in responder_upper:
            css_class = "system"
        else:
            css_class = "other"

        # Determine opacity class based on mark
        mark = getattr(fields, "mark", "")
        opacity_class = "important" if mark == "*" else "faded"

        # Start building the entry
        html_parts = [
            f'<div class="entry {css_class} {opacity_class}" id="{entry_id}">'
        ]

        # Build smart header
        entity_parts = []  # Main header line with entity info
        content_preview = ""  # Second line with content preview

        # Add task name if not root
        if task_name and task_name != "root":
            entity_parts.append(task_name)

        # Handle different responder types
        if "USER" in responder_upper:
            # Add responder with sender_entity in parens if different
            if sender_entity and sender_entity != responder:
                entity_parts.append(f"USER ({sender_entity})")
            else:
                entity_parts.append("USER")
            # Show user input preview on second line
            if content:
                preview = content.replace("\n", " ")[:60]
                if len(content) > 60:
                    preview += "..."
                content_preview = f'"{preview}"'

        elif "LLM" in responder_upper:
            # Get model info from instance - don't uppercase it
            model_label = "LLM"
            if self.model_info:
                model_label = f"LLM ({self.model_info})"

            if tool and tool_type:
                # LLM making a tool call - don't uppercase tool names
                entity_parts.append(f"{model_label} → {tool_type}[{tool}]")
            else:
                # LLM generating plain text response
                entity_parts.append(model_label)
                if content:
                    # Show first line or first 60 chars on second line
                    first_line = content.split("\n")[0].strip()
                    if first_line:
                        preview = first_line[:60]
                        if len(first_line) > 60:
                            preview += "..."
                        content_preview = f'"{preview}"'

        elif "AGENT" in responder_upper:
            # Add responder with sender_entity in parens if different
            agent_label = "AGENT"
            if sender_entity and sender_entity != responder:
                agent_label = f"AGENT ({sender_entity})"

            # Agent responding (usually tool handling)
            if tool:
                entity_parts.append(f"{agent_label}[{tool}]")
                # Show tool result preview on second line if available
                if content:
                    preview = content.replace("\n", " ")[:40]
                    if len(content) > 40:
                        preview += "..."
                    content_preview = f"→ {preview}"
            else:
                entity_parts.append(agent_label)
                if content:
                    preview = content[:50]
                    if len(content) > 50:
                        preview += "..."
                    content_preview = f'"{preview}"'

        elif "SYSTEM" in responder_upper:
            entity_parts.append("SYSTEM")
            if content:
                preview = content[:50]
                if len(content) > 50:
                    preview += "..."
                content_preview = f'"{preview}"'
        else:
            # Other responder types (like Task)
            entity_parts.append(responder)

        # Add recipient info if present
        if recipient:
            entity_parts.append(f"→ {recipient}")

        # Construct the two-line header
        header_main = " ".join(entity_parts)

        # Build the header HTML with toggle, mark, and main content on same line
        header_html = '<span class="toggle">[+]</span> '

        # Note: opacity_class already determined above

        # Add the main header content
        header_html += f'<span class="header-main">{html.escape(header_main)}</span>'

        # Add preview on second line if present
        if content_preview:
            header_html += (
                f'\n    <div class="header-content">'
                f"{html.escape(content_preview)}</div>"
            )

        # Add expandable header
        html_parts.append(
            f"""
<div class="entity-header" onclick="toggleEntry('{entry_id}')">
    {header_html}
</div>
<div id="{entry_id}_content" class="entry-content collapsed">"""
        )

        # Add collapsible sections

        # System messages (if any)
        system_content = self._extract_system_content(fields)
        if system_content:
            for idx, (label, content) in enumerate(system_content):
                section_id = f"{entry_id}_system_{idx}"
                html_parts.append(
                    self._create_collapsible_section(section_id, label, content)
                )

        # Tool information
        tool = getattr(fields, "tool", None)
        # Only add tool section if tool exists and is not empty
        if tool and tool.strip():
            tool_html = self._format_tool_section(fields, entry_id)
            html_parts.append(tool_html)

        # Main content
        content = getattr(fields, "content", "")
        if content and not (
            tool and tool.strip()
        ):  # Don't duplicate content if it's a tool
            html_parts.append(f'<div class="main-content">{html.escape(content)}</div>')

        # Metadata (recipient, blocked)
        metadata_parts = []
        recipient = getattr(fields, "recipient", None)
        if recipient:
            metadata_parts.append(f"Recipient: {recipient}")

        block = getattr(fields, "block", None)
        if block:
            metadata_parts.append(f"Blocked: {block}")

        if metadata_parts:
            html_parts.append(
                f'<div class="metadata">{" | ".join(metadata_parts)}</div>'
            )

        # Close entry content div
        html_parts.append("</div>")  # Close entry-content
        html_parts.append("</div>")  # Close entry
        return "\n".join(html_parts)

    def _extract_system_content(self, fields: BaseModel) -> List[tuple[str, str]]:
        """Extract system-related content from fields.

        Returns:
            List of (label, content) tuples
        """
        system_content = []

        # Check for common system message patterns in content
        content = getattr(fields, "content", "")
        if content:
            # Look for patterns like "[System Prompt]" or "System Reminder:"
            if "[System Prompt]" in content or "System Prompt" in content:
                system_content.append(("System Prompt", content))
            elif "[System Reminder]" in content or "System Reminder" in content:
                system_content.append(("System Reminder", content))

        return system_content

    def _create_collapsible_section(
        self, section_id: str, label: str, content: str
    ) -> str:
        """Create a collapsible section.

        Args:
            section_id: Unique ID for the section
            label: Label to display
            content: Content to show when expanded

        Returns:
            HTML string for the collapsible section
        """
        return f"""
<div class="collapsible collapsed" id="{section_id}">
    <span class="toggle" onclick="toggle('{section_id}')">[+]</span> {label}
    <div class="content">{html.escape(content)}</div>
</div>"""

    def _format_tool_section(self, fields: BaseModel, entry_id: str) -> str:
        """Format tool-related information.

        Args:
            fields: ChatDocLoggerFields containing tool information
            entry_id: Parent entry ID

        Returns:
            HTML string for the tool section
        """
        tool = getattr(fields, "tool", "")
        tool_type = getattr(fields, "tool_type", "")
        content = getattr(fields, "content", "")

        tool_id = f"{entry_id}_tool_{self.tool_counter}"
        self.tool_counter += 1

        # Try to parse content as JSON for better formatting
        try:
            if content.strip().startswith("{"):
                content_dict = json.loads(content)
                formatted_content = json.dumps(content_dict, indent=2)
                content_html = (
                    f'<pre class="code-block">{html.escape(formatted_content)}</pre>'
                )
            else:
                content_html = html.escape(content)
        except Exception:
            content_html = html.escape(content)

        # Build tool section
        tool_name = f"{tool_type}({tool})" if tool_type else tool

        return f"""
<div class="tool-section">
    <div class="collapsible collapsed" id="{tool_id}">
        <span class="toggle" onclick="toggle('{tool_id}')">[+]</span>
        <span class="tool-name">{html.escape(tool_name)}</span>
        <div class="content">{content_html}</div>
    </div>
</div>"""

    def _append_to_file(self, content: str) -> None:
        """Append content to the HTML file.

        Args:
            content: HTML content to append
        """
        try:
            with open(self.file_path, "a", encoding="utf-8") as f:
                f.write(content + "\n")
                f.flush()
        except Exception as e:
            self.logger.error(f"Failed to append to file: {e}")

    def close(self) -> None:
        """Close the HTML file with footer."""
        footer = """
    </div>
    <script>
        // Update message count
        const header = document.querySelector('.header-line div:last-child');
        if (header) {
            const messageCount = document.querySelectorAll('.entry').length;
            header.textContent = header.textContent.replace(
                /\\d+ messages/, messageCount + ' messages'
            );
        }
    </script>
</body>
</html>"""
        try:
            with open(self.file_path, "a", encoding="utf-8") as f:
                f.write(footer)
        except Exception as e:
            self.logger.error(f"Failed to write HTML footer: {e}")
</file>

<file path="langroid/utils/logging.py">
import logging
import os
import os.path
import sys
import threading
from typing import ClassVar, Dict, no_type_check

import colorlog
from rich.console import Console
from rich.markup import escape


# Define a function to set up the colored logger
def setup_colored_logging() -> None:
    # Define the log format with color codes
    log_format = "%(log_color)s%(asctime)s - %(levelname)s - %(message)s%(reset)s"
    # Create a color formatter
    color_formatter = colorlog.ColoredFormatter(
        log_format,
        datefmt="%Y-%m-%d %H:%M:%S",
        reset=True,
        log_colors={
            "DEBUG": "cyan",
            "INFO": "green",
            "WARNING": "yellow",
            "ERROR": "red",
            "CRITICAL": "red,bg_white",
        },
    )
    # Configure the root logger to use the color formatter
    handler = logging.StreamHandler()
    handler.setFormatter(color_formatter)
    logger = logging.getLogger()
    logger.addHandler(handler)
    # logger.setLevel(logging.DEBUG)


def setup_logger(
    name: str,
    level: int = logging.INFO,
    terminal: bool = False,
) -> logging.Logger:
    """
    Set up a logger of module `name` at a desired level.
    Args:
        name: module name
        level: desired logging level
    Returns:
        logger
    """
    logger = logging.getLogger(name)
    logger.setLevel(level)
    if not logger.hasHandlers() and terminal:
        handler = logging.StreamHandler()
        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
        )
        handler.setFormatter(formatter)
        logger.addHandler(handler)
    return logger


def setup_console_logger(name: str) -> logging.Logger:
    logger = setup_logger(name)
    handler = logging.StreamHandler()
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger


def setup_file_logger(
    name: str,
    filename: str,
    append: bool = False,
    log_format: bool = False,
    propagate: bool = False,
) -> logging.Logger:
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    file_mode = "a" if append else "w"
    logger = setup_logger(name, terminal=False)
    handler = logging.FileHandler(filename, mode=file_mode, encoding="utf-8")
    handler.setLevel(logging.INFO)
    if log_format:
        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
        )
    else:
        formatter = logging.Formatter("%(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.propagate = propagate
    return logger


def setup_loggers_for_package(package_name: str, level: int) -> None:
    """
    Set up loggers for all modules in a package.
    This ensures that log-levels of modules outside the package are not affected.
    Args:
        package_name: main package name
        level: desired logging level
    Returns:
    """
    import importlib
    import pkgutil

    package = importlib.import_module(package_name)
    for _, module_name, _ in pkgutil.walk_packages(
        package.__path__, package.__name__ + "."
    ):
        module = importlib.import_module(module_name)
        setup_logger(module.__name__, level)


class RichFileLogger:
    """Singleton-per-path, ref-counted, thread-safe file logger.

    • Any number of calls to `RichFileLogger(path)` yield the same object.
    • A per-instance lock guarantees that the underlying file is opened only
      once, even when many threads construct the logger concurrently.
    • A reference counter tracks how many parts of the program are using the
      logger; the FD is closed only when the counter reaches zero.
    • All writes are serialised with a dedicated write-lock.
    """

    _instances: ClassVar[Dict[str, "RichFileLogger"]] = {}
    _ref_counts: ClassVar[Dict[str, int]] = {}
    # guards _instances & _ref_counts
    _class_lock: ClassVar[threading.Lock] = threading.Lock()

    # ------------------------------------------------------------------ #
    # construction / destruction
    # ------------------------------------------------------------------ #
    def __new__(
        cls, log_file: str, append: bool = False, color: bool = True
    ) -> "RichFileLogger":
        with cls._class_lock:
            if log_file in cls._instances:
                cls._ref_counts[log_file] += 1
                return cls._instances[log_file]

            inst = super().__new__(cls)
            # create the per-instance init-lock *before* releasing class-lock
            inst._init_lock = threading.Lock()
            cls._instances[log_file] = inst
            cls._ref_counts[log_file] = 1
            return inst

    def __init__(self, log_file: str, append: bool = False, color: bool = True) -> None:
        # Double-checked locking: perform heavy init exactly once.
        if getattr(self, "_init_done", False):
            return

        if not hasattr(self, "_init_lock"):
            self._init_lock: threading.Lock = threading.Lock()

        with self._init_lock:
            if getattr(self, "_init_done", False):
                return

            os.makedirs(os.path.dirname(log_file), exist_ok=True)
            mode = "a" if append else "w"
            self._owns_file: bool = True
            try:
                self.file = open(log_file, mode, buffering=1, encoding="utf-8")
            except OSError as exc:  # EMFILE: too many open files
                if exc.errno == 24:
                    # Fallback: reuse an already-open stream to avoid creating a new FD
                    self.file = sys.stderr
                    self._owns_file = False
                else:
                    raise
            self.log_file: str = log_file
            self.color: bool = color
            self.console: Console | None = (
                Console(file=self.file, force_terminal=True, width=200)
                if color
                else None
            )
            self._write_lock = threading.Lock()
            self._init_done = True  # set last

    # ------------------------------------------------------------------ #
    # public API
    # ------------------------------------------------------------------ #
    @no_type_check
    def log(self, message: str) -> None:
        """Thread-safe write to the log file."""
        with self._write_lock:
            if self.color and self.console is not None:
                self.console.print(escape(message))
            else:
                print(message, file=self.file)
            self.file.flush()

    def close(self) -> None:
        """Decrease ref-count; close FD only when last user is done."""
        with self._class_lock:
            count = self._ref_counts.get(self.log_file, 0) - 1
            if count <= 0:
                self._ref_counts.pop(self.log_file, None)
                self._instances.pop(self.log_file, None)
                with self._write_lock:
                    if self._owns_file and not self.file.closed:
                        self.file.close()
            else:
                self._ref_counts[self.log_file] = count
</file>

<file path="langroid/utils/object_registry.py">
import time
from typing import TYPE_CHECKING, Dict, Optional, TypeAlias, TypeVar
from uuid import uuid4

from pydantic import BaseModel

if TYPE_CHECKING:
    from langroid.agent.base import Agent
    from langroid.agent.chat_agent import ChatAgent
    from langroid.agent.chat_document import ChatDocument

    # any derivative of BaseModel that has an id() method or an id attribute
    ObjWithId: TypeAlias = ChatDocument | ChatAgent | Agent
else:
    ObjWithId = BaseModel

# Define a type variable that can be any subclass of BaseModel
T = TypeVar("T", bound=BaseModel)


class ObjectRegistry:
    """A global registry to hold id -> object mappings."""

    registry: Dict[str, ObjWithId] = {}

    @classmethod
    def add(cls, obj: ObjWithId) -> str:
        """Adds an object to the registry, returning the object's ID."""
        object_id = obj.id() if callable(obj.id) else obj.id
        cls.registry[object_id] = obj
        return object_id

    @classmethod
    def get(cls, obj_id: str) -> Optional[ObjWithId]:
        """Retrieves an object by ID if it still exists."""
        return cls.registry.get(obj_id)

    @classmethod
    def register_object(cls, obj: ObjWithId) -> str:
        """Registers an object in the registry, returning the object's ID."""
        return cls.add(obj)

    @classmethod
    def remove(cls, obj_id: str) -> None:
        """Removes an object from the registry."""
        if obj_id in cls.registry:
            del cls.registry[obj_id]

    @classmethod
    def cleanup(cls) -> None:
        """Cleans up the registry by removing entries where the object is None."""
        to_remove = [key for key, value in cls.registry.items() if value is None]
        for key in to_remove:
            del cls.registry[key]

    @staticmethod
    def new_id() -> str:
        """Generates a new unique ID."""
        return str(uuid4())


def scheduled_cleanup(interval: int = 600) -> None:
    """Periodically cleans up the global registry every 'interval' seconds."""
    while True:
        ObjectRegistry.cleanup()
        time.sleep(interval)
</file>

<file path="langroid/utils/types.py">
import json
import logging
from inspect import signature
from typing import Any, Optional, Type, TypeVar, Union, get_args, get_origin

from pydantic import BaseModel

logger = logging.getLogger(__name__)
PrimitiveType = Union[int, float, bool, str]
T = TypeVar("T")


def is_instance_of(obj: Any, type_hint: Type[T] | Any) -> bool:
    """
    Check if an object is an instance of a type hint, e.g.
    to check whether x is of type `List[ToolMessage]` or type `int`
    """
    if type_hint == Any:
        return True

    if type_hint is type(obj):
        return True

    origin = get_origin(type_hint)
    args = get_args(type_hint)

    if origin is Union:
        return any(is_instance_of(obj, arg) for arg in args)

    if origin:  # e.g. List, Dict, Tuple, Set
        if isinstance(obj, origin):
            # check if all items in obj are of the required types
            if args:
                if isinstance(obj, (list, tuple, set)):
                    return all(is_instance_of(item, args[0]) for item in obj)
                if isinstance(obj, dict):
                    return all(
                        is_instance_of(k, args[0]) and is_instance_of(v, args[1])
                        for k, v in obj.items()
                    )
            return True
        else:
            return False

    return isinstance(obj, type_hint)


def to_string(msg: Any) -> str:
    """
    Best-effort conversion of arbitrary msg to str.
    Return empty string if conversion fails.
    """
    if msg is None:
        return ""
    if isinstance(msg, str):
        return msg
    if isinstance(msg, BaseModel):
        return msg.model_dump_json()
    # last resort: use json.dumps() or str() to make it a str
    try:
        return json.dumps(msg)
    except Exception:
        try:
            return str(msg)
        except Exception as e:
            logger.error(
                f"""
                Error converting msg to str: {e}", 
                """,
                exc_info=True,
            )
            return ""


def from_string(
    s: str,
    output_type: Type[PrimitiveType],
) -> Optional[PrimitiveType]:
    if output_type is int:
        try:
            return int(s)
        except ValueError:
            return None
    elif output_type is float:
        try:
            return float(s)
        except ValueError:
            return None
    elif output_type is bool:
        return s.lower() in ("true", "yes", "1")
    elif output_type is str:
        return s
    else:
        return None


def is_callable(obj: Any, k: int = 1) -> bool:
    """Check if object is callable and accepts exactly k args.

    Args:
        obj: Object to check

    Returns:
        bool: True if object is callable with k args, False otherwise
    """
    if not callable(obj):
        return False
    try:
        sig = signature(obj)
        params = list(sig.parameters.values())
        return len(params) == k
    except ValueError:
        return False
</file>

<file path="langroid/vector_store/__init__.py">
from . import base

from . import qdrantdb

from .base import VectorStoreConfig, VectorStore
from .qdrantdb import QdrantDBConfig, QdrantDB

__all__ = [
    "base",
    "VectorStore",
    "VectorStoreConfig",
    "qdrantdb",
    "QdrantDBConfig",
    "QdrantDB",
]


try:
    from . import meilisearch
    from .meilisearch import MeiliSearch, MeiliSearchConfig

    meilisearch
    MeiliSearch
    MeiliSearchConfig
    __all__.extend(["meilisearch", "MeiliSearch", "MeiliSearchConfig"])

    from . import lancedb
    from .lancedb import LanceDB, LanceDBConfig

    lancedb
    LanceDB
    LanceDBConfig
    __all__.extend(["lancedb", "LanceDB", "LanceDBConfig"])
    from . import chromadb
    from .chromadb import ChromaDBConfig, ChromaDB

    chromadb  # silence linters
    ChromaDB
    ChromaDBConfig
    __all__.extend(["chromadb", "ChromaDBConfig", "ChromaDB"])

    from . import postgres
    from .postgres import PostgresDB, PostgresDBConfig

    postgres  # silence linters
    PostgresDB
    PostgresDBConfig
    __all__.extend(["postgres", "PostgresDB", "PostgresDBConfig"])

    from . import weaviatedb
    from .weaviatedb import WeaviateDBConfig, WeaviateDB

    weaviatedb
    WeaviateDB
    WeaviateDBConfig
    __all__.extend(["weaviatedb", "WeaviateDB", "WeaviateDBConfig"])

    from . import pineconedb
    from .pineconedb import PineconeDB, PineconeDBConfig

    pineconedb
    PineconeDB
    PineconeDBConfig
    __all__.extend(["pineconedb", "PineconeDB", "PineconeDBConfig"])
except ImportError:
    pass
</file>

<file path="langroid/vector_store/chromadb.py">
import json
import logging
from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple

from langroid.embedding_models.base import (
    EmbeddingModelsConfig,
)
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.exceptions import LangroidImportError
from langroid.mytypes import Document
from langroid.utils.configuration import settings
from langroid.utils.output.printing import print_long_text
from langroid.vector_store.base import VectorStore, VectorStoreConfig

logger = logging.getLogger(__name__)


class ChromaDBConfig(VectorStoreConfig):
    collection_name: str = "temp"
    storage_path: str = ".chroma/data"
    distance: Literal["cosine", "l2", "ip"] = "cosine"
    construction_ef: int = 100
    search_ef: int = 100
    max_neighbors: int = 16
    embedding: EmbeddingModelsConfig = OpenAIEmbeddingsConfig()
    host: str = "127.0.0.1"
    port: int = 6333


class ChromaDB(VectorStore):
    def __init__(self, config: ChromaDBConfig = ChromaDBConfig()):
        super().__init__(config)
        try:
            import chromadb
        except ImportError:
            raise LangroidImportError("chromadb", "chromadb")
        self.config = config
        self.client = chromadb.Client(
            chromadb.config.Settings(
                # chroma_db_impl="duckdb+parquet",
                # is_persistent=bool(config.storage_path),
                persist_directory=config.storage_path,
            )
        )
        if self.config.collection_name is not None:
            self.create_collection(
                self.config.collection_name,
                replace=self.config.replace_collection,
            )

    def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
        """Clear all collections in the vector store with the given prefix."""

        if not really:
            logger.warning("Not deleting all collections, set really=True to confirm")
            return 0
        coll = [c for c in self.client.list_collections() if c.name.startswith(prefix)]
        if len(coll) == 0:
            logger.warning(f"No collections found with prefix {prefix}")
            return 0
        n_empty_deletes = 0
        n_non_empty_deletes = 0
        for c in coll:
            n_empty_deletes += c.count() == 0
            n_non_empty_deletes += c.count() > 0
            self.client.delete_collection(name=c.name)
        logger.warning(
            f"""
            Deleted {n_empty_deletes} empty collections and
            {n_non_empty_deletes} non-empty collections.
            """
        )
        return n_empty_deletes + n_non_empty_deletes

    def clear_empty_collections(self) -> int:
        colls = self.client.list_collections()
        n_deletes = 0
        for coll in colls:
            if coll.count() == 0:
                n_deletes += 1
                self.client.delete_collection(name=coll.name)
        return n_deletes

    def list_collections(self, empty: bool = False) -> List[str]:
        """
        List non-empty collections in the vector store.
        Args:
            empty (bool, optional): Whether to list empty collections.
        Returns:
            List[str]: List of non-empty collection names.
        """
        colls = self.client.list_collections()
        if empty:
            return [coll.name for coll in colls]
        return [coll.name for coll in colls if coll.count() > 0]

    def create_collection(self, collection_name: str, replace: bool = False) -> None:
        """
        Create a collection in the vector store, optionally replacing an existing
            collection if `replace` is True.
        Args:
            collection_name (str): Name of the collection to create or replace.
            replace (bool, optional): Whether to replace an existing collection.
                Defaults to False.

        """
        self.config.collection_name = collection_name
        if collection_name in self.list_collections(empty=True) and replace:
            logger.warning(f"Replacing existing collection {collection_name}")
            self.client.delete_collection(collection_name)
        self.collection = self.client.create_collection(
            name=self.config.collection_name,
            embedding_function=self.embedding_fn,
            get_or_create=not replace,
            metadata={
                "hnsw:space": self.config.distance,
                "hnsw:construction_ef": self.config.construction_ef,
                "hnsw:search_ef": self.config.search_ef,
                # we could expose other configs, see:
                # https://docs.trychroma.com/docs/collections/configure
            },
        )

    def add_documents(self, documents: Sequence[Document]) -> None:
        super().maybe_add_ids(documents)
        if documents is None:
            return
        contents: List[str] = [document.content for document in documents]
        # convert metadatas to dicts so chroma can handle them
        metadata_dicts: List[dict[str, Any]] = [
            d.metadata.dict_bool_int() 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]

        colls = self.list_collections(empty=True)
        if self.config.collection_name is None:
            raise ValueError("No collection name set, cannot ingest docs")
        if self.config.collection_name not in colls:
            self.create_collection(self.config.collection_name, replace=True)

        self.collection.add(
            # embedding_models=embedding_models,
            documents=contents,
            metadatas=metadata_dicts,
            ids=ids,
        )

    def get_all_documents(self, where: str = "") -> List[Document]:
        filter = json.loads(where) if where else None
        results = self.collection.get(
            include=["documents", "metadatas"],
            where=filter,
        )
        results["documents"] = [results["documents"]]
        results["metadatas"] = [results["metadatas"]]
        return self._docs_from_results(results)

    def get_documents_by_ids(self, ids: List[str]) -> List[Document]:
        # get them one by one since chroma mangles the order of the results
        # when fetched from a list of ids.
        results = [
            self.collection.get(ids=[id], include=["documents", "metadatas"])
            for id in ids
        ]
        final_results = {}
        final_results["documents"] = [[r["documents"][0] for r in results]]
        final_results["metadatas"] = [[r["metadatas"][0] for r in results]]
        return self._docs_from_results(final_results)

    def delete_collection(self, collection_name: str) -> None:
        try:
            self.client.delete_collection(name=collection_name)
        except Exception:
            pass

    def similar_texts_with_scores(
        self, text: str, k: int = 1, where: Optional[str] = None
    ) -> List[Tuple[Document, float]]:
        n = self.collection.count()
        filter = json.loads(where) if where else None
        results = self.collection.query(
            query_texts=[text],
            n_results=min(n, k),
            where=filter,
            include=["documents", "distances", "metadatas"],
        )
        docs = self._docs_from_results(results)
        # 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]:
        """
        Helper function to convert results from ChromaDB to a list of Documents
        Args:
            results (dict): results from ChromaDB

        Returns:
            List[Document]: list of Documents
        """
        if len(results["documents"][0]) == 0:
            return []
        contents = results["documents"][0]
        if settings.debug:
            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]
            if m["window_ids"].strip() == "":
                m["window_ids"] = []
            else:
                m["window_ids"] = m["window_ids"].split(",")
        docs = [
            self.config.document_class(
                content=d, metadata=self.config.metadata_class(**m)
            )
            for d, m in zip(contents, metadatas)
        ]
        return docs
</file>

<file path="langroid/vector_store/lancedb.py">
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Generator,
    List,
    Optional,
    Sequence,
    Tuple,
    Type,
)

import pandas as pd
from dotenv import load_dotenv
from pydantic import BaseModel, ValidationError, create_model

if TYPE_CHECKING:
    from lancedb.query import LanceVectorQueryBuilder

from langroid.embedding_models.base import (
    EmbeddingModelsConfig,
)
from langroid.embedding_models.models import OpenAIEmbeddingsConfig
from langroid.exceptions import LangroidImportError
from langroid.mytypes import Document
from langroid.utils.configuration import settings
from langroid.utils.pydantic_utils import (
    dataframe_to_document_model,
    dataframe_to_documents,
)
from langroid.vector_store.base import VectorStore, VectorStoreConfig

try:
    import lancedb
    from lancedb.pydantic import LanceModel, Vector

    has_lancedb = True
except ImportError:
    has_lancedb = False

logger = logging.getLogger(__name__)


class LanceDBConfig(VectorStoreConfig):
    cloud: bool = False
    collection_name: str | None = "temp"
    storage_path: str = ".lancedb/data"
    embedding: EmbeddingModelsConfig = OpenAIEmbeddingsConfig()
    distance: str = "cosine"


class LanceDB(VectorStore):
    def __init__(self, config: LanceDBConfig = LanceDBConfig()):
        super().__init__(config)
        if not has_lancedb:
            raise LangroidImportError("lancedb", "lancedb")

        self.config: LanceDBConfig = config
        self.host = config.host
        self.port = config.port
        self.is_from_dataframe = False  # were docs ingested from a dataframe?
        self.df_metadata_columns: List[str] = []  # metadata columns from dataframe

        load_dotenv()
        if self.config.cloud:
            logger.warning(
                "LanceDB Cloud is not available yet. Switching to local storage."
            )
            config.cloud = False
        else:
            try:
                self.client = lancedb.connect(
                    uri=config.storage_path,
                )
            except Exception as e:
                new_storage_path = config.storage_path + ".new"
                logger.warning(
                    f"""
                    Error connecting to local LanceDB at {config.storage_path}:
                    {e}
                    Switching to {new_storage_path}
                    """
                )
                self.client = lancedb.connect(
                    uri=new_storage_path,
                )

    def clear_empty_collections(self) -> int:
        coll_names = self.list_collections()
        n_deletes = 0
        for name in coll_names:
            nr = self.client.open_table(name).head(1).shape[0]
            if nr == 0:
                n_deletes += 1
                self.client.drop_table(name)
        return n_deletes

    def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
        """Clear all collections with the given prefix."""
        if not really:
            logger.warning("Not deleting all collections, set really=True to confirm")
            return 0
        coll_names = [
            c for c in self.list_collections(empty=True) if c.startswith(prefix)
        ]
        if len(coll_names) == 0:
            logger.warning(f"No collections found with prefix {prefix}")
            return 0
        n_empty_deletes = 0
        n_non_empty_deletes = 0
        for name in coll_names:
            nr = self.client.open_table(name).head(1).shape[0]
            n_empty_deletes += nr == 0
            n_non_empty_deletes += nr > 0
            self.client.drop_table(name)
        logger.warning(
            f"""
            Deleted {n_empty_deletes} empty collections and
            {n_non_empty_deletes} non-empty collections.
            """
        )
        return n_empty_deletes + n_non_empty_deletes

    def list_collections(self, empty: bool = False) -> List[str]:
        """
        Returns:
            List of collection names that have at least one vector.

        Args:
            empty (bool, optional): Whether to include empty collections.
        """
        colls = self.client.table_names(limit=None)
        if len(colls) == 0:
            return []
        if empty:  # include empty tbls
            return colls  # type: ignore
        counts = [self.client.open_table(coll).head(1).shape[0] for coll in colls]
        return [coll for coll, count in zip(colls, counts) if count > 0]

    def _create_lance_schema(self, doc_cls: Type[Document]) -> Type[BaseModel]:
        """
        NOTE: NOT USED, but leaving it here as it may be useful.

        Create a subclass of LanceModel with fields:
         - id (str)
         - Vector field that has dims equal to
            the embedding dimension of the embedding model, and a data field of type
            DocClass.
         - other fields from doc_cls

        Args:
            doc_cls