elna / main.py
David Chu
doc: improve code comments
d38794d unverified
raw
history blame
1.71 kB
import streamlit as st
from google import genai
from app import agent, config
def format_output(response: list[dict]) -> tuple[str, str]:
try:
answer = ""
citations = {}
for statement in response:
text = statement["text"].strip()
answer = (
answer + f"\n{text}"
if text.startswith("*") or text.startswith("-")
else answer + f" {text}"
)
citation_ids = []
for source in statement.get("sources", []):
source_str = f"[{source['title']}]({source['url']})"
if not (citation_id := citations.get(source_str)):
citation_id = len(citations) + 1
citations[source_str] = citation_id
citation_ids.append(citation_id)
if citation_ids:
answer += " ".join(f"[^{i}]" for i in sorted(citation_ids))
except KeyError as err:
print(err)
return str(response), ""
footnotes = "\n".join(f"[^{id}]: {citation}" for citation, id in citations.items())
return answer, footnotes
def main():
gemini_client = genai.Client(api_key=config.settings.google_api_key)
st.title("Elna")
with st.form("search", border=False):
query = st.text_input("Your medical question")
submit = st.form_submit_button("Ask")
response = st.empty()
if submit:
with st.spinner("Thinking...", show_time=True):
output = agent.respond(gemini_client, query)
answer, footnotes = format_output(output)
response.markdown(f"{answer}\n\n{footnotes}")
if __name__ == "__main__":
main()