feat: adding rest of the project files.

.gitignore

@@ -0,0 +1,179 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+pages*/
+*milvus*
+
+.gradio

.python-version

@@ -0,0 +1 @@
+3.12

pyproject.toml

@@ -0,0 +1,39 @@
+[project]
+name = "insurance-assistants"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [
+    { name = "shamik", email = "[email protected]" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "beautifulsoup4>=4.13.4",
+    "colpali-engine>=0.3.10",
+    "crawlee[all]>=0.6.9",
+    "duckduckgo-search>=8.0.2",
+    "gradio-pdf>=0.0.22",
+    "huggingface-hub[hf-xet]>=0.32.0",
+    "llama-index>=0.12.37",
+    "llama-index-llms-huggingface-api>=0.4.3",
+    "llama-index-tools-playwright>=0.2.0",
+    "markdownify>=1.1.0",
+    "pdf2image>=1.17.0",
+    "pydantic>=2.11.5",
+    "pymilvus>=2.5.10",
+    "pyprojroot>=0.3.0",
+    "python-dotenv>=1.1.0",
+    "qwen-vl-utils>=0.0.11",
+    "requests>=2.32.3",
+    "smolagents[gradio]>=1.16.1",
+    "torchvision>=0.21.0",
+    "wikipedia-api>=0.8.1",
+]
+
+
+[dependency-groups]
+dev = [
+    "jupyter>=1.1.1",
+    "pyrefly>=0.16.2",
+    "ruff>=0.11.11",
+]

requirements.txt

@@ -0,0 +1,656 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements.txt
+accelerate==1.7.0
+    # via peft
+aiofiles==24.1.0
+    # via gradio
+aiohappyeyeballs==2.6.1
+    # via aiohttp
+aiohttp==3.12.4
+    # via llama-index-core
+aiosignal==1.3.2
+    # via aiohttp
+aiosqlite==0.21.0
+    # via llama-index-core
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.9.0
+    # via
+    #   gradio
+    #   httpx
+    #   openai
+    #   starlette
+apify-fingerprint-datapoints==0.0.2
+    # via crawlee
+arrow==1.3.0
+    # via cookiecutter
+attrs==25.3.0
+    # via aiohttp
+av==14.4.0
+    # via qwen-vl-utils
+banks==2.1.2
+    # via llama-index-core
+beautifulsoup4==4.13.4
+    # via
+    #   insurance-assistants (pyproject.toml)
+    #   crawlee
+    #   llama-index-readers-file
+    #   llama-index-tools-playwright
+    #   markdownify
+binaryornot==0.4.4
+    # via cookiecutter
+blessed==1.21.0
+    # via inquirer
+brotli==1.1.0
+    # via httpx
+browserforge==1.2.3
+    # via crawlee
+cachetools==6.0.0
+    # via crawlee
+certifi==2025.1.31
+    # via
+    #   crawlee
+    #   curl-cffi
+    #   httpcore
+    #   httpx
+    #   llama-cloud
+    #   requests
+cffi==1.17.1
+    # via curl-cffi
+chardet==5.2.0
+    # via binaryornot
+charset-normalizer==3.4.2
+    # via requests
+click==8.2.1
+    # via
+    #   browserforge
+    #   cookiecutter
+    #   duckduckgo-search
+    #   llama-cloud-services
+    #   nltk
+    #   typer
+    #   uvicorn
+colorama==0.4.6
+    # via
+    #   crawlee
+    #   griffe
+colpali-engine==0.3.10
+    # via insurance-assistants (pyproject.toml)
+cookiecutter==2.6.0
+    # via crawlee
+crawlee==0.6.9
+    # via insurance-assistants (pyproject.toml)
+cssselect==1.3.0
+    # via parsel
+curl-cffi==0.11.1
+    # via crawlee
+dataclasses-json==0.6.7
+    # via llama-index-core
+deprecated==1.2.18
+    # via
+    #   banks
+    #   llama-index-core
+dirtyjson==1.0.8
+    # via llama-index-core
+distro==1.9.0
+    # via openai
+duckduckgo-search==8.0.2
+    # via insurance-assistants (pyproject.toml)
+editor==1.6.6
+    # via inquirer
+eval-type-backport==0.2.2
+    # via crawlee
+fastapi==0.115.12
+    # via gradio
+ffmpy==0.5.0
+    # via gradio
+filelock==3.18.0
+    # via
+    #   huggingface-hub
+    #   tldextract
+    #   torch
+    #   transformers
+filetype==1.2.0
+    # via llama-index-core
+frozenlist==1.6.0
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2025.5.1
+    # via
+    #   gradio-client
+    #   huggingface-hub
+    #   llama-index-core
+    #   torch
+gradio==5.31.0
+    # via
+    #   gradio-pdf
+    #   smolagents
+gradio-client==1.10.1
+    # via gradio
+gradio-pdf==0.0.22
+    # via insurance-assistants (pyproject.toml)
+greenlet==3.2.2
+    # via
+    #   playwright
+    #   sqlalchemy
+griffe==1.7.3
+    # via banks
+groovy==0.1.2
+    # via gradio
+grpcio==1.67.1
+    # via pymilvus
+h11==0.16.0
+    # via
+    #   httpcore
+    #   uvicorn
+h2==4.2.0
+    # via httpx
+hf-xet==1.1.2
+    # via huggingface-hub
+hpack==4.1.0
+    # via h2
+html5lib==1.1
+    # via crawlee
+httpcore==1.0.9
+    # via httpx
+httpx==0.28.1
+    # via
+    #   crawlee
+    #   gradio
+    #   gradio-client
+    #   llama-cloud
+    #   llama-index-core
+    #   openai
+    #   safehttpx
+huggingface-hub==0.32.3
+    # via
+    #   insurance-assistants (pyproject.toml)
+    #   accelerate
+    #   gradio
+    #   gradio-client
+    #   llama-index-llms-huggingface-api
+    #   peft
+    #   smolagents
+    #   tokenizers
+    #   transformers
+hyperframe==6.1.0
+    # via h2
+idna==3.10
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+    #   tldextract
+    #   yarl
+inquirer==3.4.0
+    # via crawlee
+jaro-winkler==2.0.3
+    # via crawlee
+jinja2==3.1.6
+    # via
+    #   banks
+    #   cookiecutter
+    #   gradio
+    #   smolagents
+    #   torch
+jiter==0.10.0
+    # via openai
+jmespath==1.0.1
+    # via parsel
+joblib==1.5.1
+    # via
+    #   nltk
+    #   scikit-learn
+llama-cloud==0.1.23
+    # via
+    #   llama-cloud-services
+    #   llama-index-indices-managed-llama-cloud
+llama-cloud-services==0.6.25
+    # via llama-parse
+llama-index==0.12.38
+    # via insurance-assistants (pyproject.toml)
+llama-index-agent-openai==0.4.8
+    # via
+    #   llama-index
+    #   llama-index-program-openai
+llama-index-cli==0.4.1
+    # via llama-index
+llama-index-core==0.12.38
+    # via
+    #   llama-cloud-services
+    #   llama-index
+    #   llama-index-agent-openai
+    #   llama-index-cli
+    #   llama-index-embeddings-openai
+    #   llama-index-indices-managed-llama-cloud
+    #   llama-index-llms-huggingface-api
+    #   llama-index-llms-openai
+    #   llama-index-multi-modal-llms-openai
+    #   llama-index-program-openai
+    #   llama-index-question-gen-openai
+    #   llama-index-readers-file
+    #   llama-index-readers-llama-parse
+    #   llama-index-tools-playwright
+llama-index-embeddings-openai==0.3.1
+    # via
+    #   llama-index
+    #   llama-index-cli
+llama-index-indices-managed-llama-cloud==0.7.1
+    # via llama-index
+llama-index-llms-huggingface-api==0.5.0
+    # via insurance-assistants (pyproject.toml)
+llama-index-llms-openai==0.3.44
+    # via
+    #   llama-index
+    #   llama-index-agent-openai
+    #   llama-index-cli
+    #   llama-index-multi-modal-llms-openai
+    #   llama-index-program-openai
+    #   llama-index-question-gen-openai
+llama-index-multi-modal-llms-openai==0.4.3
+    # via llama-index
+llama-index-program-openai==0.3.1
+    # via
+    #   llama-index
+    #   llama-index-question-gen-openai
+llama-index-question-gen-openai==0.3.0
+    # via llama-index
+llama-index-readers-file==0.4.8
+    # via llama-index
+llama-index-readers-llama-parse==0.4.0
+    # via llama-index
+llama-index-tools-playwright==0.2.0
+    # via insurance-assistants (pyproject.toml)
+llama-parse==0.6.25
+    # via llama-index-readers-llama-parse
+lxml==5.4.0
+    # via
+    #   beautifulsoup4
+    #   duckduckgo-search
+    #   parsel
+markdown-it-py==3.0.0
+    # via rich
+markdownify==1.1.0
+    # via insurance-assistants (pyproject.toml)
+markupsafe==3.0.2
+    # via
+    #   gradio
+    #   jinja2
+marshmallow==3.26.1
+    # via dataclasses-json
+mdurl==0.1.2
+    # via markdown-it-py
+milvus-lite==2.4.12
+    # via pymilvus
+more-itertools==10.7.0
+    # via crawlee
+mpmath==1.3.0
+    # via sympy
+multidict==6.4.4
+    # via
+    #   aiohttp
+    #   yarl
+mypy-extensions==1.1.0
+    # via typing-inspect
+nest-asyncio==1.6.0
+    # via llama-index-core
+networkx==3.5
+    # via
+    #   llama-index-core
+    #   torch
+nltk==3.9.1
+    # via
+    #   llama-index
+    #   llama-index-core
+numpy==2.2.6
+    # via
+    #   accelerate
+    #   colpali-engine
+    #   gradio
+    #   llama-index-core
+    #   pandas
+    #   peft
+    #   scikit-learn
+    #   scipy
+    #   torchvision
+    #   transformers
+nvidia-cublas-cu12==12.4.5.8
+    # via
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cuda-cupti-cu12==12.4.127
+    # via torch
+nvidia-cuda-nvrtc-cu12==12.4.127
+    # via torch
+nvidia-cuda-runtime-cu12==12.4.127
+    # via torch
+nvidia-cudnn-cu12==9.1.0.70
+    # via torch
+nvidia-cufft-cu12==11.2.1.3
+    # via torch
+nvidia-curand-cu12==10.3.5.147
+    # via torch
+nvidia-cusolver-cu12==11.6.1.9
+    # via torch
+nvidia-cusparse-cu12==12.3.1.170
+    # via
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cusparselt-cu12==0.6.2
+    # via torch
+nvidia-nccl-cu12==2.21.5
+    # via torch
+nvidia-nvjitlink-cu12==12.4.127
+    # via
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
+    #   torch
+nvidia-nvtx-cu12==12.4.127
+    # via torch
+openai==1.82.1
+    # via
+    #   llama-index-agent-openai
+    #   llama-index-embeddings-openai
+    #   llama-index-llms-openai
+orjson==3.10.18
+    # via gradio
+packaging==25.0
+    # via
+    #   accelerate
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   marshmallow
+    #   parsel
+    #   peft
+    #   qwen-vl-utils
+    #   transformers
+pandas==2.2.3
+    # via
+    #   gradio
+    #   llama-index-readers-file
+    #   pymilvus
+parsel==1.10.0
+    # via crawlee
+pdf2image==1.17.0
+    # via insurance-assistants (pyproject.toml)
+peft==0.15.2
+    # via colpali-engine
+pillow==11.2.1
+    # via
+    #   colpali-engine
+    #   gradio
+    #   llama-index-core
+    #   pdf2image
+    #   qwen-vl-utils
+    #   smolagents
+    #   torchvision
+platformdirs==4.3.8
+    # via
+    #   banks
+    #   llama-cloud-services
+playwright==1.52.0
+    # via
+    #   crawlee
+    #   llama-index-tools-playwright
+primp==0.15.0
+    # via duckduckgo-search
+propcache==0.3.1
+    # via
+    #   aiohttp
+    #   yarl
+protego==0.4.0
+    # via crawlee
+protobuf==6.31.1
+    # via pymilvus
+psutil==7.0.0
+    # via
+    #   accelerate
+    #   crawlee
+    #   peft
+pycparser==2.22
+    # via cffi
+pydantic==2.11.5
+    # via
+    #   insurance-assistants (pyproject.toml)
+    #   banks
+    #   crawlee
+    #   fastapi
+    #   gradio
+    #   llama-cloud
+    #   llama-cloud-services
+    #   llama-index-core
+    #   openai
+    #   pydantic-settings
+pydantic-core==2.33.2
+    # via pydantic
+pydantic-settings==2.6.1
+    # via crawlee
+pydub==0.25.1
+    # via gradio
+pyee==13.0.0
+    # via
+    #   crawlee
+    #   playwright
+pygments==2.19.1
+    # via rich
+pymilvus==2.5.10
+    # via insurance-assistants (pyproject.toml)
+pypdf==5.5.0
+    # via llama-index-readers-file
+pyprojroot==0.3.0
+    # via insurance-assistants (pyproject.toml)
+python-dateutil==2.9.0.post0
+    # via
+    #   arrow
+    #   pandas
+python-dotenv==1.1.0
+    # via
+    #   insurance-assistants (pyproject.toml)
+    #   llama-cloud-services
+    #   pydantic-settings
+    #   pymilvus
+    #   smolagents
+python-multipart==0.0.20
+    # via gradio
+python-slugify==8.0.4
+    # via cookiecutter
+pytz==2025.2
+    # via pandas
+pyyaml==6.0.2
+    # via
+    #   accelerate
+    #   cookiecutter
+    #   gradio
+    #   huggingface-hub
+    #   llama-index-core
+    #   peft
+    #   transformers
+qwen-vl-utils==0.0.11
+    # via insurance-assistants (pyproject.toml)
+readchar==4.2.1
+    # via inquirer
+regex==2024.11.6
+    # via
+    #   nltk
+    #   tiktoken
+    #   transformers
+requests==2.32.3
+    # via
+    #   insurance-assistants (pyproject.toml)
+    #   colpali-engine
+    #   cookiecutter
+    #   huggingface-hub
+    #   llama-index-core
+    #   qwen-vl-utils
+    #   requests-file
+    #   smolagents
+    #   tiktoken
+    #   tldextract
+    #   transformers
+    #   wikipedia-api
+requests-file==2.1.0
+    # via tldextract
+rich==14.0.0
+    # via
+    #   cookiecutter
+    #   smolagents
+    #   typer
+ruff==0.11.12
+    # via gradio
+runs==1.2.2
+    # via editor
+safehttpx==0.1.6
+    # via gradio
+safetensors==0.5.3
+    # via
+    #   accelerate
+    #   peft
+    #   transformers
+scikit-learn==1.6.1
+    # via crawlee
+scipy==1.15.3
+    # via
+    #   colpali-engine
+    #   scikit-learn
+semantic-version==2.10.0
+    # via gradio
+setuptools==80.9.0
+    # via
+    #   pymilvus
+    #   torch
+shellingham==1.5.4
+    # via typer
+six==1.17.0
+    # via
+    #   html5lib
+    #   markdownify
+    #   python-dateutil
+smolagents==1.17.0
+    # via insurance-assistants (pyproject.toml)
+sniffio==1.3.1
+    # via
+    #   anyio
+    #   openai
+sortedcollections==2.1.0
+    # via crawlee
+sortedcontainers==2.4.0
+    # via
+    #   crawlee
+    #   sortedcollections
+soupsieve==2.7
+    # via beautifulsoup4
+sqlalchemy==2.0.41
+    # via llama-index-core
+starlette==0.46.2
+    # via
+    #   fastapi
+    #   gradio
+striprtf==0.0.26
+    # via llama-index-readers-file
+sympy==1.13.1
+    # via torch
+tenacity==9.1.2
+    # via llama-index-core
+text-unidecode==1.3
+    # via python-slugify
+threadpoolctl==3.6.0
+    # via scikit-learn
+tiktoken==0.9.0
+    # via llama-index-core
+tldextract==5.3.0
+    # via crawlee
+tokenizers==0.21.1
+    # via transformers
+tomlkit==0.13.2
+    # via gradio
+torch==2.6.0
+    # via
+    #   accelerate
+    #   colpali-engine
+    #   peft
+    #   torchvision
+torchvision==0.21.0
+    # via insurance-assistants (pyproject.toml)
+tqdm==4.67.1
+    # via
+    #   huggingface-hub
+    #   llama-index-core
+    #   milvus-lite
+    #   nltk
+    #   openai
+    #   peft
+    #   transformers
+transformers==4.51.3
+    # via
+    #   colpali-engine
+    #   peft
+triton==3.2.0
+    # via torch
+typer==0.16.0
+    # via
+    #   crawlee
+    #   gradio
+types-python-dateutil==2.9.0.20250516
+    # via arrow
+typing-extensions==4.13.2
+    # via
+    #   aiosqlite
+    #   anyio
+    #   beautifulsoup4
+    #   crawlee
+    #   fastapi
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   llama-index-core
+    #   openai
+    #   pydantic
+    #   pydantic-core
+    #   pyee
+    #   pyprojroot
+    #   sqlalchemy
+    #   torch
+    #   typer
+    #   typing-inspect
+    #   typing-inspection
+typing-inspect==0.9.0
+    # via
+    #   dataclasses-json
+    #   llama-index-core
+typing-inspection==0.4.1
+    # via pydantic
+tzdata==2025.2
+    # via pandas
+ujson==5.10.0
+    # via pymilvus
+urllib3==2.4.0
+    # via requests
+uvicorn==0.34.2
+    # via gradio
+w3lib==2.3.1
+    # via parsel
+wcwidth==0.2.13
+    # via blessed
+webencodings==0.5.1
+    # via html5lib
+websockets==15.0.1
+    # via gradio-client
+wikipedia-api==0.8.1
+    # via insurance-assistants (pyproject.toml)
+wrapt==1.17.2
+    # via
+    #   deprecated
+    #   llama-index-core
+xmod==1.8.1
+    # via
+    #   editor
+    #   runs
+yarl==1.20.0
+    # via
+    #   aiohttp
+    #   crawlee
+zstandard==0.23.0
+    # via httpx

src/insurance_assistants/__init__.py

@@ -0,0 +1,2 @@
+def hello() -> str:
+    return "Hello from insurance-assistants!"

src/insurance_assistants/agents.py

@@ -0,0 +1,131 @@
+from pathlib import Path
+
+from dotenv import find_dotenv, load_dotenv
+from huggingface_hub import InferenceClient
+from smolagents import (
+    CodeAgent,
+    DuckDuckGoSearchTool,
+    FinalAnswerTool,
+    InferenceClientModel,
+    PythonInterpreterTool,
+    Tool,
+    ToolCallingAgent,
+    VisitWebpageTool,
+    WikipediaSearchTool,
+)
+
+from src.insurance_assistants.complex_rag import RAG
+from src.insurance_assistants.consts import PROMPT_PREFIX
+
+_ = load_dotenv(dotenv_path=find_dotenv())
+rag_app = RAG()
+# FIXME Comment the following if you want to reprocess everything
+rag_app.vectordb_id = "policy_wordings"
+
+
+class InsuranceInfoRetriever(Tool):
+    name = "InsuranceInfoRetriever"
+    description = "Retrieves information from insurance documents."
+    inputs = {
+        "query": {"type": "string", "description": "The query to search for."},
+    }
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        client = InferenceClient(
+            provider="hyperbolic",
+            bill_to="VitalNest",
+        )
+        results = rag_app.search_documents(query)
+        img_paths = [Path(res[0]) for res in results]
+
+        grouped_images = [rag_app.encode_image_to_base64(pth) for pth in img_paths]
+        chat_template = [
+            {
+                "role": "system",
+                "content": """You find answers from the relevant documents. Answer only 
+        from these documents. If answer isn't available return 'Question cannot be answered based
+        on the documents provided.' """,
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{image}"},
+                    }
+                    for image in grouped_images
+                ]
+                + [{"type": "text", "text": query}],
+            },
+        ]
+        completion = client.chat.completions.create(
+            model="Qwen/Qwen2.5-VL-72B-Instruct",
+            messages=chat_template,
+            temperature=0.1,
+            max_tokens=10_000,
+        )
+        answer = completion.choices[0].message.content
+        if answer:
+            answer += f"The information was retrived from the following documents: {img_paths}"
+        return answer if answer else ""
+
+
+insurance_agent = CodeAgent(
+    tools=[InsuranceInfoRetriever(), FinalAnswerTool()],
+    model=InferenceClientModel(bill_to="VitalNest", temperature=0.1),
+    additional_authorized_imports=["os", "requests", "bs4", "pil", "base64", "io"],
+    max_steps=1,
+    verbosity_level=-1,
+    name="insurance_agent",
+    description="You answer health insurance questions based on the InsuranceInfoRetriever "
+    "tool. All health insurance questions must be answered by you.",
+)
+websearch_agent = ToolCallingAgent(
+    model=InferenceClientModel(
+        model_id="Qwen/Qwen3-30B-A3B", bill_to="VitalNest", temperature=0.1
+    ),
+    tools=[
+        VisitWebpageTool(max_output_length=20_000),
+        DuckDuckGoSearchTool(max_results=5),
+        FinalAnswerTool(),
+    ],
+    max_steps=4,
+    verbosity_level=-1,
+    name="web_search_agent",
+    planning_interval=2,
+    description="Searches the web with a particular query.",
+)
+
+wikipedia_agent = ToolCallingAgent(
+    model=InferenceClientModel(
+        model_id="Qwen/Qwen3-30B-A3B", bill_to="VitalNest", temperature=0.1
+    ),
+    tools=[
+        WikipediaSearchTool(user_agent="WikiAssistant ([email protected])"),
+        FinalAnswerTool(),
+    ],
+    max_steps=3,
+    verbosity_level=-1,
+    name="wikipedia_agent",
+    description="Searches Wikipedia for a topic.",
+)
+
+manager_agent = CodeAgent(
+    tools=[FinalAnswerTool(), PythonInterpreterTool()],
+    additional_authorized_imports=["os"],
+    model=InferenceClientModel(
+        model_id="Qwen/Qwen3-235B-A22B",
+        bill_to="VitalNest",
+        temperature=0.1,
+    ),
+    managed_agents=[websearch_agent, wikipedia_agent, insurance_agent],
+    max_steps=10,
+    planning_interval=2,
+    verbosity_level=-1,
+    add_base_tools=True,
+    name="Versatile_Multi_Agent",
+    description="Answer health insurance related questions from pre-defined set of "
+    "health insurance documents, search wikipedia and the web for general information.",
+)
+manager_agent.system_prompt = manager_agent.system_prompt + PROMPT_PREFIX

src/insurance_assistants/complex_rag.py

@@ -0,0 +1,600 @@
+# %%
+import base64
+import concurrent.futures
+import logging
+import os
+import shutil
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import torch
+from colpali_engine.models import (
+    # ColPali,
+    # ColPaliProcessor,
+    ColQwen2_5,
+    ColQwen2_5_Processor,
+)
+from colpali_engine.utils.torch_utils import ListDataset, get_torch_device
+from dotenv import find_dotenv, load_dotenv
+from openai import OpenAI
+from pdf2image import convert_from_path
+from PIL import Image
+from pymilvus import DataType, MilvusClient
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from src.insurance_assistants.consts import PROJECT_ROOT_DIR, PROMPT
+
+# Setup logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# %%
+model_name = "vidore/colqwen2.5-v0.2"
+# model_name = "vidore/colpali-v1.2"
+device = get_torch_device()
+
+model = ColQwen2_5.from_pretrained(
+    # model = ColPali.from_pretrained(
+    pretrained_model_name_or_path=model_name,
+    torch_dtype=torch.bfloat16,
+    device_map=device,
+).eval()
+
+processor = ColQwen2_5_Processor.from_pretrained(
+    # processor = ColPaliProcessor.from_pretrained(
+    pretrained_model_name_or_path=model_name,
+    use_fast=True,
+)
+_ = load_dotenv(dotenv_path=find_dotenv(raise_error_if_not_found=True))
+openai_client = OpenAI()
+
+
+# %%
+class MilvusManager:
+    def __init__(self, milvus_uri, collection_name, create_collection, dim=128):
+        """
+        Initializes the MilvusManager.
+
+        Args:
+            milvus_uri (str): URI for Milvus server.
+            collection_name (str): Name of the collection.
+            create_collection (bool): Whether to create a new collection.
+            dim (int, optional): Dimension of the vector. Defaults to 128.
+        """
+        self.client = MilvusClient(uri=milvus_uri)
+        self.collection_name = collection_name
+        if self.client.has_collection(collection_name=self.collection_name):
+            self.client.load_collection(collection_name)
+        self.dim = dim
+        self.max_doc_id = 0
+
+        if create_collection:
+            self.create_collection()
+            self.create_index()
+
+    def create_collection(self):
+        """
+        Creates a new collection in Milvus. Drops existing collection if present.
+        """
+        if self.client.has_collection(collection_name=self.collection_name):
+            self.client.drop_collection(collection_name=self.collection_name)
+        schema = self.client.create_schema(
+            auto_id=True,
+            enable_dynamic_fields=True,
+        )
+        schema.add_field(field_name="pk", datatype=DataType.INT64, is_primary=True)
+        schema.add_field(
+            field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self.dim
+        )
+        schema.add_field(field_name="seq_id", datatype=DataType.INT16)
+        schema.add_field(field_name="doc_id", datatype=DataType.INT64)
+        schema.add_field(field_name="doc", datatype=DataType.VARCHAR, max_length=65535)
+
+        self.client.create_collection(
+            collection_name=self.collection_name, schema=schema
+        )
+
+    def create_index(self):
+        """
+        Creates a vector index for the collection in Milvus.
+        """
+        self.client.release_collection(collection_name=self.collection_name)
+        self.client.drop_index(
+            collection_name=self.collection_name, index_name="vector"
+        )
+        index_params = self.client.prepare_index_params()
+        index_params.add_index(
+            field_name="vector",
+            index_name="vector_index",
+            index_type="FLAT",
+            metric_type="IP",
+            params={
+                "M": 16,
+                "efConstruction": 500,
+            },
+        )
+
+        self.client.create_index(
+            collection_name=self.collection_name, index_params=index_params, sync=True
+        )
+
+    def create_scalar_index(self):
+        """
+        Creates a scalar index for the doc_id field in Milvus.
+        """
+        self.client.release_collection(collection_name=self.collection_name)
+
+        index_params = self.client.prepare_index_params()
+        index_params.add_index(
+            field_name="doc_id",
+            index_name="int32_index",
+            index_type="INVERTED",
+        )
+
+        self.client.create_index(
+            collection_name=self.collection_name, index_params=index_params, sync=True
+        )
+
+    def search(self, data, topk):
+        """
+        Searches for the top-k most similar documents in Milvus.
+
+        Args:
+            data (np.ndarray): Query vector.
+            topk (int): Number of top results to return.
+
+        Returns:
+            list: List of (score, doc_id) tuples.
+        """
+        search_params = {"metric_type": "IP", "params": {}}
+        results = self.client.search(
+            self.collection_name,
+            data,
+            limit=50,
+            output_fields=["vector", "seq_id", "doc_id"],
+            search_params=search_params,
+        )
+        doc_ids = set()
+        for r_id in range(len(results)):
+            for r in range(len(results[r_id])):
+                doc_ids.add(results[r_id][r]["entity"]["doc_id"])
+
+        scores = []
+
+        def rerank_single_doc(doc_id, data, client, collection_name):
+            doc_colqwen_vecs = client.query(
+                collection_name=collection_name,
+                filter=f"doc_id in [{doc_id}, {doc_id + 1}]",
+                output_fields=["seq_id", "vector", "doc"],
+                limit=1000,
+            )
+            doc_vecs = np.vstack(
+                [doc_colqwen_vecs[i]["vector"] for i in range(len(doc_colqwen_vecs))]
+            )
+            score = np.dot(data, doc_vecs.T).max(1).sum()
+            return (score, doc_id)
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=300) as executor:
+            futures = {
+                executor.submit(
+                    rerank_single_doc, doc_id, data, self.client, self.collection_name
+                ): doc_id
+                for doc_id in doc_ids
+            }
+            for future in concurrent.futures.as_completed(futures):
+                score, doc_id = future.result()
+                scores.append((score, doc_id))
+
+        scores.sort(key=lambda x: x[0], reverse=True)
+        if len(scores) >= topk:
+            return scores[:topk]
+        else:
+            return scores
+
+    def insert(self, data):
+        """
+        Inserts a document's vectors and metadata into Milvus.
+
+        Args:
+            data (dict): Dictionary containing 'colqwen_vecs', 'doc_id', and 'filepath'.
+        """
+        colqwen_vecs = [vec for vec in data["colqwen_vecs"]]
+        seq_length = len(colqwen_vecs)
+        doc_ids = [data["doc_id"] for _ in range(seq_length)]
+        seq_ids = list(range(seq_length))
+        docs = [data["filepath"] for _ in range(seq_length)]
+        # docs = [""] * seq_length
+        # docs[0] = data["filepath"]
+
+        self.client.insert(
+            self.collection_name,
+            [
+                {
+                    "vector": colqwen_vecs[i],
+                    "seq_id": seq_ids[i],
+                    "doc_id": doc_ids[i],
+                    "doc": docs[i],
+                }
+                for i in range(seq_length)
+            ],
+        )
+
+    def get_images_as_doc(self, images_with_vectors: list):
+        """
+        Converts a list of image vectors and filepaths into Milvus insertable format.
+
+        Args:
+            images_with_vectors (list): List of dicts with 'colqwen_vecs' and 'filepath'.
+
+        Returns:
+            list: List of dicts ready for Milvus insertion.
+        """
+        images_data = []
+
+        for i in range(len(images_with_vectors)):
+            self.max_doc_id += 1
+            data = {
+                "colqwen_vecs": images_with_vectors[i]["colqwen_vecs"],
+                "doc_id": self.max_doc_id,
+                "filepath": images_with_vectors[i]["filepath"],
+            }
+            images_data.append(data)
+
+        return images_data
+
+    def insert_images_data(self, image_data):
+        """
+        Inserts multiple images' data into Milvus.
+
+        Args:
+            image_data (list): List of image data dicts.
+        """
+        data = self.get_images_as_doc(image_data)
+
+        for i in range(len(data)):
+            self.insert(data[i])
+
+
+# %%
+class VectorProcessor:
+    def __init__(
+        self,
+        id: str,
+        create_collection=True,
+    ):
+        """
+        Initializes the VectorProcessor with Milvus, Colqwen, and PDF managers.
+
+        Args:
+            id (str): Unique identifier for the session/user.
+            create_collection (bool, optional): Whether to create a new collection. Defaults to True.
+        """
+        # hashed_id = hashlib.md5(id.encode()).hexdigest()[:8]
+        # milvus_db_name = f"milvus_{hashed_id}.db"
+        milvus_db_name = (
+            PROJECT_ROOT_DIR / f"src/insurance_assistants/milvus_{id}.db"
+        ).as_posix()
+        self.milvus_manager = MilvusManager(milvus_db_name, f"{id}", create_collection)
+        self.colqwen_manager = ColqwenManager()
+        self.pdf_manager = PdfManager()
+
+    def index(
+        self,
+        pdf_path: str,
+        id: str,
+        max_pages: int,
+    ):
+        """
+        Indexes a PDF file by converting pages to images, embedding them, and storing in Milvus.
+
+        Args:
+            pdf_path (str): Path to the PDF file.
+            id (str): Unique identifier.
+            max_pages (int): Maximum number of pages to process.
+
+        Returns:
+            list: List of saved image paths.
+        """
+        logger.info(f"Indexing {pdf_path}, id: {id}, max_pages: {max_pages}")
+
+        image_paths = self.pdf_manager.save_images(id, pdf_path, max_pages)
+
+        logger.info(f"Saved {len(image_paths)} images")
+
+        colqwen_vecs = self.colqwen_manager.process_images(image_paths)
+
+        images_data = [
+            {"colqwen_vecs": colqwen_vecs[i], "filepath": image_paths[i]}
+            for i in range(len(image_paths))
+        ]
+
+        logger.info(f"Inserting {len(images_data)} images data to Milvus")
+
+        self.milvus_manager.insert_images_data(images_data)
+
+        logger.info("Indexing completed")
+
+        return image_paths
+
+    def search(self, search_queries: list[str]):
+        logger.info(f"Searching for {len(search_queries)} queries")
+
+        final_res = []
+
+        for query in search_queries:
+            logger.info(f"Searching for query: {query}")
+            query_vec = self.colqwen_manager.process_text([query])[0]
+            search_res = self.milvus_manager.search(query_vec, topk=4)
+            logger.info(f"Search result: {search_res} for query: {query}")
+            final_res.append(search_res)
+
+        return final_res
+
+
+# %%
+class PdfManager:
+    def __init__(self):
+        """
+        Initializes the PdfManager.
+        """
+        pass
+
+    def clear_and_recreate_dir(self, output_folder):
+        logger.info(f"Clearing output folder {output_folder}")
+
+        if os.path.exists(output_folder):
+            shutil.rmtree(output_folder)
+
+        os.makedirs(output_folder)
+
+    def save_images(
+        self, id, pdf_path, max_pages, pages: list[int] = None, output_folder=None
+    ) -> list[str]:
+        """
+        Saves images of PDF pages to disk.
+
+        Args:
+            id (str): Unique identifier.
+            pdf_path (str): Path to the PDF file.
+            max_pages (int): Maximum number of pages to save.
+            pages (list[int], optional): Specific pages to save. Defaults to None.
+
+        Returns:
+            list[str]: List of saved image file paths.
+        """
+        output_folder = (
+            Path(output_folder) if output_folder is not None else output_folder
+        )
+        if output_folder is None:
+            output_folder = PROJECT_ROOT_DIR / f"src/insurance_assistants/pages/{id}/"
+        if not output_folder.exists():
+            output_folder.mkdir(parents=True, exist_ok=True)
+        images = convert_from_path(pdf_path=pdf_path)
+
+        logger.info(
+            f"Saving images from {pdf_path} to {output_folder}. Max pages: {max_pages}"
+        )
+
+        # self.clear_and_recreate_dir(output_folder)
+
+        num_page_processed = 0
+
+        for i, image in enumerate(images):
+            if max_pages and num_page_processed >= max_pages:
+                break
+
+            if pages and i not in pages:
+                continue
+
+            full_save_path = output_folder / f"{id}_page_{i + 1}.png"
+
+            # logger.debug(f"Saving image to {full_save_path}")
+
+            image.save(fp=full_save_path, format="PNG")
+
+            num_page_processed += 1
+
+        return [
+            f"{output_folder}/{id}_page_{i + 1}.png" for i in range(num_page_processed)
+        ]
+
+
+# %%
+class ColqwenManager:
+    def get_images(self, paths: list[str]) -> list[Image.Image]:
+        """
+        Loads images from file paths.
+
+        Args:
+            paths (list[str]): List of image file paths.
+
+        Returns:
+            list[Image.Image]: List of PIL Image objects.
+        """
+        return [Image.open(path) for path in paths]
+
+    def process_images(self, image_paths: list[str], batch_size=5):
+        logger.info(f"Processing {len(image_paths)} image_paths")
+
+        images = self.get_images(image_paths)
+
+        dataloader = DataLoader(
+            dataset=ListDataset[str](images),
+            batch_size=batch_size,
+            shuffle=False,
+            collate_fn=lambda x: processor.process_images(x),
+        )
+
+        ds: list[torch.Tensor] = []
+        for batch_doc in tqdm(dataloader):
+            with torch.no_grad():
+                batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
+                embeddings_doc = model(**batch_doc)
+            ds.extend(list(torch.unbind(embeddings_doc.to(device))))
+
+        ds_np = [d.float().cpu().numpy() for d in ds]
+
+        return ds_np
+
+    def process_text(self, texts: list[str]):
+        logger.info(f"Processing {len(texts)} texts")
+
+        dataloader = DataLoader(
+            dataset=ListDataset[str](texts),
+            batch_size=1,
+            shuffle=False,
+            collate_fn=lambda x: processor.process_queries(x),
+        )
+
+        qs: list[torch.Tensor] = []
+        for batch_query in dataloader:
+            with torch.no_grad():
+                batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
+                embeddings_query = model(**batch_query)
+
+            qs.extend(list(torch.unbind(embeddings_query.to(device))))
+
+        qs_np = [q.float().cpu().numpy() for q in qs]
+
+        return qs_np
+
+
+# %%
+# def generate_uuid(state):
+#     """
+#     Generates or retrieves a UUID for the user session.
+
+#     Args:
+#         state (dict): State dictionary containing 'user_uuid'.
+
+#     Returns:
+#         str: UUID string.
+#     """
+#     # Check if UUID already exists in session state
+#     if state["user_uuid"] is None:
+#         # Generate a new UUID if not already set
+#         state["user_uuid"] = str(uuid.uuid4())
+
+#     return state["user_uuid"]
+
+
+class RAG:
+    def __init__(self):
+        """
+        Initializes the RAG.
+        """
+        self.vectordb_id = None
+        self.img_path_dir = PROJECT_ROOT_DIR / "src/insurance_assistants/pages/"
+
+    def create_vector_db(
+        self,
+        vectordb_id="policy_wordings",
+        dir=PROJECT_ROOT_DIR / "data",
+        max_pages=200,
+    ):
+        """
+        Uploads a PDF file, converts it to images, and indexes it.
+
+        Args:
+            state (dict): State dictionary for user session.
+            file: Uploaded file object.
+            max_pages (int, optional): Maximum number of pages to process. Defaults to 100.
+
+        Returns:
+            str: Status message.
+        """
+
+        logger.info(f"Converting files in: {dir}.")
+
+        try:
+            for idx, f in enumerate((dir / "policy_wordings").iterdir()):
+                if idx == 0:
+                    vectorprocessor = VectorProcessor(
+                        id=vectordb_id, create_collection=True
+                    )
+                    self.vectordb_id = vectordb_id
+                _ = vectorprocessor.index(pdf_path=f, id=f.stem, max_pages=max_pages)
+                vectorprocessor.milvus_manager.client.close()
+            return f"✅ Created the vector_db: milvus_{vectordb_id} under `src` dir."
+        except Exception as err:
+            return f"❌ Error creating vector_db: {err}"
+
+    def search_documents(self, query):
+        if self.vectordb_id is None:
+            raise Exception(
+                "Create the vector db first by invoking `create_vector_db`."
+            )
+        try:
+            vectorprocessor = VectorProcessor(
+                id=self.vectordb_id, create_collection=False
+            )
+
+            search_results = vectorprocessor.search(search_queries=[query])[0]
+
+            check_res = vectorprocessor.milvus_manager.client.query(
+                collection_name=self.vectordb_id,
+                filter=f"doc_id in {[d[1] for d in search_results]}",
+                output_fields=["doc_id", "doc"],
+            )
+            vectorprocessor.milvus_manager.client.close()
+            img_path_doc_id = set((i["doc"], i["doc_id"]) for i in check_res)
+
+            logger.info("✅ Retrieved the images for answering query.")
+            return img_path_doc_id
+
+        except Exception as err:
+            return f"❌ Error during search: {err}"
+
+    def encode_image_to_base64(self, image_path):
+        """
+        Encodes an image file to a base64 string.
+
+        Args:
+            image_path (str): Path to the image file.
+
+        Returns:
+            str: Base64 encoded image string.
+        """
+        image = Image.open(image_path)
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    def query_gpt4o_mini(self, query, image_path):
+        """
+        Queries the OpenAI GPT-4o-mini model with a query and images.
+
+        Args:
+            query (str): The user query.
+            image_path (list): List of image file paths.
+
+        Returns:
+            str: The AI response.
+        """
+        try:
+            base64_images = [self.encode_image_to_base64(pth) for pth in image_path]
+
+            response = openai_client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": PROMPT.format(query=query)}
+                        ]
+                        + [
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": f"data:image/jpeg;base64,{im}"},
+                            }
+                            for im in base64_images
+                        ],
+                    }
+                ],
+                max_tokens=500,
+            )
+            return response.choices[0].message.content
+        except Exception as err:
+            return f"Unable to generate the final output due to: {err}."

src/insurance_assistants/consts.py

@@ -0,0 +1,33 @@
+from pyprojroot import find_root, has_file
+
+PROJECT_ROOT_DIR = find_root(criterion=has_file("README.md"))
+
+PROMPT = """
+You are a smart assistant designed to answer questions about a PDF document.
+You are given relevant information in the form of PDF pages. Use them to construct a short response to the question, and cite your sources (page numbers, etc).
+If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
+Give detailed and extensive answers, only containing info in the pages you are given.
+You can answer using information contained in plots and figures if necessary.
+Answer in the same language as the query.
+
+Query: {query}
+PDF pages:
+"""
+
+PRIMARY_HEADING = """# An Agentic RAG for Health Insurance Documents
+#### This agent answers health insurance related questions from pre-ingested set \
+of health insurance documents, search wikipedia, search the web \
+and execute basic python code. \
+
+The pre-ingested health insurance documents can be viewed under `PDF Viewer` section.
+"""
+
+PROMPT_PREFIX = """\n
+Apart from all the above instructions that we have given to you, FOLLOW the Additional Instructions below:
+```
+For any health insurance related queries, always use the `insurance_agent` first and return the results.
+You are allowed to rephrase any query and detail it if required. When in doubt always ask the user a follow up question.
+Don't assume anything.
+```
+
+"""

src/insurance_assistants/ui.py

@@ -0,0 +1,337 @@
+import logging
+import mimetypes
+import os
+import re
+import shutil
+
+import gradio as gr
+from dotenv import load_dotenv
+from gradio_pdf import PDF
+from smolagents.gradio_ui import _process_action_step, _process_final_answer_step
+from smolagents.memory import ActionStep, FinalAnswerStep, MemoryStep, PlanningStep
+from smolagents.models import ChatMessageStreamDelta
+
+# from smolagents import CodeAgent, InferenceClientModel
+from src.insurance_assistants.agents import manager_agent
+from src.insurance_assistants.consts import PRIMARY_HEADING, PROJECT_ROOT_DIR
+
+load_dotenv(override=True)
+
+# Setup logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class UI:
+    """A one-line interface to launch your agent in Gradio"""
+
+    def __init__(self, file_upload_folder: str | None = None):
+        self.file_upload_folder = file_upload_folder
+        if self.file_upload_folder is not None:
+            if not os.path.exists(file_upload_folder):
+                os.mkdir(file_upload_folder)
+
+    def pull_messages_from_step(
+        self, step_log: MemoryStep, skip_model_outputs: bool = False
+    ):
+        """Extract ChatMessage objects from agent steps with proper nesting.
+
+        Args:
+            step_log: The step log to display as gr.ChatMessage objects.
+            skip_model_outputs: If True, skip the model outputs when creating the gr.ChatMessage objects:
+                This is used for instance when streaming model outputs have already been displayed.
+        """
+        if isinstance(step_log, ActionStep):
+            yield from _process_action_step(step_log, skip_model_outputs)
+        elif isinstance(step_log, PlanningStep):
+            pass
+        #     yield from _process_planning_step(step_log, skip_model_outputs)
+        elif isinstance(step_log, FinalAnswerStep):
+            yield from _process_final_answer_step(step_log)
+        else:
+            raise ValueError(f"Unsupported step type: {type(step_log)}")
+
+    def stream_to_gradio(
+        self,
+        agent,
+        task: str,
+        task_images: list | None = None,
+        reset_agent_memory: bool = False,
+        additional_args: dict | None = None,
+    ):
+        """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
+        intermediate_text = ""
+        for step_log in agent.run(
+            task,
+            images=task_images,
+            stream=True,
+            reset=reset_agent_memory,
+            additional_args=additional_args,
+        ):
+            # Track tokens if model provides them
+            if getattr(agent.model, "last_input_token_count", None) is not None:
+                if isinstance(step_log, (ActionStep, PlanningStep)):
+                    step_log.input_token_count = agent.model.last_input_token_count
+                    step_log.output_token_count = agent.model.last_output_token_count
+
+            if isinstance(step_log, MemoryStep):
+                intermediate_text = ""
+                for message in self.pull_messages_from_step(
+                    step_log,
+                    # If we're streaming model outputs, no need to display them twice
+                    skip_model_outputs=getattr(agent, "stream_outputs", False),
+                ):
+                    yield message
+            elif isinstance(step_log, ChatMessageStreamDelta):
+                intermediate_text += step_log.content or ""
+                yield intermediate_text
+
+    def interact_with_agent(self, prompt, messages, session_state):
+        # Get or create session-specific agent
+        if "agent" not in session_state:
+            # session_state["agent"] = CodeAgent(tools=[], model=InfenceClientModel())
+            session_state["agent"] = manager_agent
+        # Adding monitoring
+        try:
+            # log the existence of agent memory
+            has_memory = hasattr(session_state["agent"], "memory")
+            logger.info(f"Agent has memory: {has_memory}")
+            if has_memory:
+                logger.info(f"Memory type: {type(session_state['agent'].memory)}")
+
+            messages.append(gr.ChatMessage(role="user", content=prompt))
+            yield messages
+
+            for msg in self.stream_to_gradio(
+                agent=session_state["agent"],
+                task=prompt,
+                reset_agent_memory=False,
+            ):
+                messages.append(msg)
+                yield messages
+            yield messages
+        except Exception as e:
+            logger.info(f"Error in interaction: {str(e)}")
+            raise
+
+    def upload_file(
+        self,
+        file,
+        file_uploads_log,
+        allowed_file_types=[
+            "application/pdf",
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+            "text/plain",
+        ],
+    ):
+        """
+        Handle file uploads, default allowed types are .pdf, .docx, and .txt
+        """
+        if file is None:
+            return gr.Textbox("No file uploaded", visible=True), file_uploads_log
+
+        try:
+            mime_type, _ = mimetypes.guess_type(file.name)
+        except Exception as e:
+            return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log
+
+        if mime_type not in allowed_file_types:
+            return gr.Textbox("File type disallowed", visible=True), file_uploads_log
+
+        # Sanitize file name
+        original_name = os.path.basename(file.name)
+        sanitized_name = re.sub(
+            r"[^\w\-.]", "_", original_name
+        )  # Replace any non-alphanumeric, non-dash, or non-dot characters with underscores
+
+        type_to_ext = {}
+        for ext, t in mimetypes.types_map.items():
+            if t not in type_to_ext:
+                type_to_ext[t] = ext
+
+        # Ensure the extension correlates to the mime type
+        sanitized_name = sanitized_name.split(".")[:-1]
+        sanitized_name.append("" + type_to_ext[mime_type])
+        sanitized_name = "".join(sanitized_name)
+
+        # Save the uploaded file to the specified folder
+        file_path = os.path.join(
+            self.file_upload_folder, os.path.basename(sanitized_name)
+        )
+        shutil.copy(file.name, file_path)
+
+        return gr.Textbox(
+            f"File uploaded: {file_path}", visible=True
+        ), file_uploads_log + [file_path]
+
+    def log_user_message(self, text_input, file_uploads_log):
+        return (
+            text_input
+            + (
+                f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"
+                if len(file_uploads_log) > 0
+                else ""
+            ),
+            gr.Textbox(
+                value="",
+                interactive=False,
+                placeholder="Please wait while the agent answers your question",
+            ),
+            gr.Button(interactive=False),
+        )
+
+    def list_pdfs(self, dir=PROJECT_ROOT_DIR / "data/policy_wordings"):
+        file_names = [f.name for f in dir.iterdir()]
+        return file_names
+
+    def interrupt_agent(self, session_state):
+        if "agent" not in session_state:
+            session_state["agent"] = manager_agent
+        agent = session_state["agent"]
+        agent.interrupt()
+        return
+
+    def display_pdf(self, pdf_selector):
+        return PDF(
+            value=(f"{PROJECT_ROOT_DIR}/data/policy_wordings/{pdf_selector}"),
+            label="PDF Viewer",
+            show_label=True,
+        )
+
+    def launch(self, **kwargs):
+        with gr.Blocks(fill_height=True) as demo:
+            gr.Markdown(value=PRIMARY_HEADING)
+
+            @gr.render()
+            def layout(request: gr.Request):
+                # Render layout with sidebar
+                with gr.Blocks(
+                    fill_height=True,
+                ):
+                    file_uploads_log = gr.State([])
+                    with gr.Sidebar():
+                        gr.Markdown(
+                            value="""#### <span style="color:red"> The `interrupt` button doesn't stop the process instantaneously.</span>
+                                    <span style="color:green">You can continue to use the application upon pressing the interrupt button.</span>
+"""
+                        )
+                        with gr.Group():
+                            gr.Markdown(
+                                value="**Your question, please...**", container=True
+                            )
+                            text_input = gr.Textbox(
+                                lines=3,
+                                label="Your question, please...",
+                                container=False,
+                                placeholder="Enter your prompt here and press Shift+Enter or press `Run`",
+                            )
+                            run_btn = gr.Button(value="Run", variant="primary")
+                            agent_interrup_btn = gr.Button(
+                                value="Interrupt", variant="stop"
+                            )
+
+                        # If an upload folder is provided, enable the upload feature
+                        if self.file_upload_folder is not None:
+                            upload_file = gr.File(label="Upload a file")
+                            upload_status = gr.Textbox(
+                                label="Upload Status",
+                                interactive=False,
+                                visible=False,
+                            )
+                            upload_file.change(
+                                fn=self.upload_file,
+                                inputs=[upload_file, file_uploads_log],
+                                outputs=[upload_status, file_uploads_log],
+                            )
+
+                        gr.HTML("<br><br><h4><center>Powered by:</center></h4>")
+                        with gr.Row():
+                            gr.HTML("""<div style="display: flex; align-items: center; gap: 8px; font-family: system-ui, -apple-system, sans-serif;">
+                    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png" style="width: 32px; height: 32px; object-fit: contain;" alt="logo">
+                    <a target="_blank" href="https://github.com/huggingface/smolagents"><b>huggingface/smolagents</b></a>
+                    </div>""")
+
+                    # Add session state to store session-specific data
+                    session_state = gr.State({})
+                    # Initialize empty state for each session
+                    stored_messages = gr.State([])
+                    chatbot = gr.Chatbot(
+                        label="Health Insurance Agent",
+                        type="messages",
+                        avatar_images=(
+                            None,
+                            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
+                        ),
+                        resizeable=False,
+                        scale=1,
+                        elem_id="Insurance-Agent",
+                    )
+                    with gr.Group():
+                        gr.Markdown("### 📈 PDF Viewer")
+                        pdf_choices = self.list_pdfs()
+                        pdf_selector = gr.Dropdown(
+                            choices=pdf_choices,
+                            label="Select a PDF",
+                            info="Choose one",
+                            show_label=True,
+                            interactive=True,
+                        )
+                        pdf_viewer = PDF(
+                            label="PDF Viewer",
+                            show_label=True,
+                        )
+                        pdf_selector.change(
+                            fn=self.display_pdf, inputs=pdf_selector, outputs=pdf_viewer
+                        )
+
+                    text_input.submit(
+                        fn=self.log_user_message,
+                        inputs=[text_input, file_uploads_log],
+                        outputs=[stored_messages, text_input, run_btn],
+                    ).then(
+                        fn=self.interact_with_agent,
+                        # Include session_state in function calls
+                        inputs=[stored_messages, chatbot, session_state],
+                        outputs=[chatbot],
+                    ).then(
+                        fn=lambda: (
+                            gr.Textbox(
+                                interactive=True,
+                                placeholder="Enter your prompt here or press `Run`",
+                            ),
+                            gr.Button(interactive=True),
+                        ),
+                        inputs=None,
+                        outputs=[text_input, run_btn],
+                    )
+                    run_btn.click(
+                        fn=self.log_user_message,
+                        inputs=[text_input, file_uploads_log],
+                        outputs=[stored_messages, text_input, run_btn],
+                    ).then(
+                        fn=self.interact_with_agent,
+                        # Include session_state in function calls
+                        inputs=[stored_messages, chatbot, session_state],
+                        outputs=[chatbot],
+                    ).then(
+                        fn=lambda: (
+                            gr.Textbox(
+                                interactive=True,
+                                placeholder="Enter your prompt here or press `Run`",
+                            ),
+                            gr.Button(interactive=True),
+                        ),
+                        inputs=None,
+                        outputs=[text_input, run_btn],
+                    )
+                    agent_interrup_btn.click(
+                        fn=self.interrupt_agent,
+                        inputs=[session_state],
+                    )
+
+        demo.launch(debug=True, **kwargs)
+
+
+# if __name__=="__main__":
+#     UI().launch(share=True,
+# allowed_paths=[(PROJECT_ROOT_DIR /"data/policy_wordings").as_posix()])