changing
Browse files
app.py
CHANGED
@@ -1,13 +1,14 @@
|
|
1 |
import os
|
2 |
import streamlit as st
|
3 |
from together import Together
|
4 |
-
from
|
5 |
-
from
|
6 |
|
7 |
# --- Configuration ---
|
|
|
8 |
TOGETHER_API_KEY = os.environ.get("pilotikval")
|
9 |
if not TOGETHER_API_KEY:
|
10 |
-
st.error("Missing
|
11 |
st.stop()
|
12 |
|
13 |
# Initialize TogetherAI client
|
@@ -51,64 +52,71 @@ vectorstore = Chroma(
|
|
51 |
persist_directory=persist_directory,
|
52 |
embedding_function=embeddings
|
53 |
)
|
54 |
-
retriever = vectorstore.as_retriever(search_kwargs={"k":20})
|
55 |
|
56 |
-
# System prompt template
|
57 |
def build_system(context: str) -> dict:
|
58 |
-
return {
|
59 |
-
"
|
60 |
-
"
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
|
|
|
|
|
|
|
|
65 |
|
66 |
# Initialize chat history
|
67 |
if 'chat_history' not in st.session_state:
|
68 |
-
st.session_state.chat_history = []
|
69 |
|
70 |
-
# Get user input
|
71 |
-
|
72 |
|
73 |
-
# Tabs for
|
74 |
chat_tab, clear_tab = st.tabs(["Chat", "Clear History"])
|
75 |
|
76 |
with chat_tab:
|
77 |
-
# Display
|
78 |
for msg in st.session_state.chat_history:
|
79 |
st.chat_message(msg['role']).write(msg['content'])
|
80 |
|
81 |
-
#
|
82 |
-
if
|
83 |
-
#
|
84 |
-
st.chat_message("user").write(
|
85 |
-
st.session_state.chat_history.append({"role": "user", "content":
|
86 |
|
87 |
-
# Retrieve
|
88 |
-
docs = retriever.get_relevant_documents(
|
89 |
context = "\n---\n".join([d.page_content for d in docs])
|
90 |
|
91 |
-
#
|
92 |
messages = [build_system(context)]
|
93 |
-
# Include full history
|
94 |
for m in st.session_state.chat_history:
|
95 |
-
|
96 |
-
messages.append(m)
|
97 |
|
98 |
-
#
|
99 |
response_container = st.chat_message("assistant")
|
100 |
-
|
101 |
answer = ""
|
|
|
|
|
102 |
for token in client.chat.completions.create(
|
103 |
model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
104 |
messages=messages,
|
|
|
|
|
105 |
stream=True
|
106 |
):
|
107 |
-
if hasattr(token, 'choices'):
|
108 |
delta = token.choices[0].delta.content
|
109 |
answer += delta
|
110 |
-
|
111 |
-
|
|
|
112 |
st.session_state.chat_history.append({"role": "assistant", "content": answer})
|
113 |
|
114 |
with clear_tab:
|
@@ -116,5 +124,5 @@ with clear_tab:
|
|
116 |
st.session_state.chat_history = []
|
117 |
st.experimental_rerun()
|
118 |
|
119 |
-
#
|
120 |
# vectorstore.persist()
|
|
|
1 |
import os
|
2 |
import streamlit as st
|
3 |
from together import Together
|
4 |
+
from langchain_community.vectorstores import Chroma
|
5 |
+
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
|
6 |
|
7 |
# --- Configuration ---
|
8 |
+
# TogetherAI API key (env var name pilotikval)
|
9 |
TOGETHER_API_KEY = os.environ.get("pilotikval")
|
10 |
if not TOGETHER_API_KEY:
|
11 |
+
st.error("Missing pilotikval environment variable.")
|
12 |
st.stop()
|
13 |
|
14 |
# Initialize TogetherAI client
|
|
|
52 |
persist_directory=persist_directory,
|
53 |
embedding_function=embeddings
|
54 |
)
|
55 |
+
retriever = vectorstore.as_retriever(search_kwargs={"k": 20}) # k=20
|
56 |
|
57 |
+
# System prompt template with instruction for detailed long answers
|
58 |
def build_system(context: str) -> dict:
|
59 |
+
return {
|
60 |
+
"role": "system",
|
61 |
+
"content": (
|
62 |
+
"You are an expert medical assistant. Provide a thorough, detailed, and complete answer. "
|
63 |
+
"If you don't know, say you don't know.\n"
|
64 |
+
"Use the following context from medical docs to answer.\n\n"
|
65 |
+
"Context:\n" + context
|
66 |
+
)
|
67 |
+
}
|
68 |
+
|
69 |
+
st.title("🩺 DocChatter RAG (Streaming & Memory)")
|
70 |
|
71 |
# Initialize chat history
|
72 |
if 'chat_history' not in st.session_state:
|
73 |
+
st.session_state.chat_history = [] # list of dicts {role, content}
|
74 |
|
75 |
+
# Get user input at top level
|
76 |
+
user_prompt = st.chat_input("Ask anything about your docs…")
|
77 |
|
78 |
+
# Tabs for UI
|
79 |
chat_tab, clear_tab = st.tabs(["Chat", "Clear History"])
|
80 |
|
81 |
with chat_tab:
|
82 |
+
# Display existing chat
|
83 |
for msg in st.session_state.chat_history:
|
84 |
st.chat_message(msg['role']).write(msg['content'])
|
85 |
|
86 |
+
# On new input
|
87 |
+
if user_prompt:
|
88 |
+
# Echo user
|
89 |
+
st.chat_message("user").write(user_prompt)
|
90 |
+
st.session_state.chat_history.append({"role": "user", "content": user_prompt})
|
91 |
|
92 |
+
# Retrieve top-k docs
|
93 |
+
docs = retriever.get_relevant_documents(user_prompt)
|
94 |
context = "\n---\n".join([d.page_content for d in docs])
|
95 |
|
96 |
+
# Build message sequence: system + full history
|
97 |
messages = [build_system(context)]
|
|
|
98 |
for m in st.session_state.chat_history:
|
99 |
+
messages.append(m)
|
|
|
100 |
|
101 |
+
# Prepare streaming response
|
102 |
response_container = st.chat_message("assistant")
|
103 |
+
stream_placeholder = response_container.empty()
|
104 |
answer = ""
|
105 |
+
|
106 |
+
# Stream tokens
|
107 |
for token in client.chat.completions.create(
|
108 |
model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
109 |
messages=messages,
|
110 |
+
max_tokens=22048,
|
111 |
+
temperature=0.1,
|
112 |
stream=True
|
113 |
):
|
114 |
+
if hasattr(token, 'choices') and token.choices[0].delta.content:
|
115 |
delta = token.choices[0].delta.content
|
116 |
answer += delta
|
117 |
+
stream_placeholder.write(answer)
|
118 |
+
|
119 |
+
# Save assistant response
|
120 |
st.session_state.chat_history.append({"role": "assistant", "content": answer})
|
121 |
|
122 |
with clear_tab:
|
|
|
124 |
st.session_state.chat_history = []
|
125 |
st.experimental_rerun()
|
126 |
|
127 |
+
# (Optional) persist new docs
|
128 |
# vectorstore.persist()
|