src/legisqa_local/components/display.py

"""Display components for LegisQA"""

import streamlit as st
from legisqa_local.utils.text import escape_markdown, replace_legis_ids_with_urls
from legisqa_local.utils.usage import display_api_usage
from legisqa_local.utils.formatting import render_retrieved_chunks
from legisqa_local.config.models import PROVIDER_MODELS


def render_example_queries():
    """Render example queries in an expander"""
    with st.expander("Example Queries"):
        st.write(
            """

```
What are the themes around artificial intelligence?
```

```
Write a well cited 3 paragraph essay on food insecurity.
```

```
Create a table summarizing major climate change ideas with columns legis_id, title, idea.
```

```
Write an action plan to keep social security solvent.
```

```
Suggest reforms that would benefit the Medicaid program.
```

        """
        )


def render_response(
    response: dict,
    model_info: dict,
    provider: str,
    should_escape_markdown: bool,
    should_add_legis_urls: bool,
    tag: str | None = None,
):
    """Render a RAG response with usage information and retrieved chunks"""
    response_text = response["aimessage"].content
    if should_escape_markdown:
        response_text = escape_markdown(response_text)
    if should_add_legis_urls:
        response_text = replace_legis_ids_with_urls(response_text)

    with st.container(border=True):
        if tag is None:
            st.write("Response")
        else:
            st.write(f"Response ({tag})")
        st.info(response_text)

    display_api_usage(response["aimessage"], model_info, provider, tag=tag)
    render_retrieved_chunks(response["docs"], tag=tag)