-
Notifications
You must be signed in to change notification settings - Fork 4
Expand file tree
/
Copy pathapp.py
More file actions
454 lines (387 loc) Β· 18.1 KB
/
app.py
File metadata and controls
454 lines (387 loc) Β· 18.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
from __future__ import annotations
import concurrent.futures
import os
from pathlib import Path
from typing import Dict, List, Any
import json
import time
import base64
from datetime import datetime
import ast
import streamlit as st
from agent_builder import build_local_agent
from ingestion import process_documents
from retrieval import retrieve_documents
from think_tank import ThinkTank
from utils import export_meeting, clean_name
from agno.memory.v2 import UserMemory
DB_FILE = Path("projects_db copy.json")
TEMPLATE_FILE = Path("scientist_templates.json")
def img_to_base64(path):
return base64.b64encode(Path(path).read_bytes()).decode()
img_b64 = img_to_base64("assets/Logo_tau.png")
if "rows" not in st.session_state:
st.session_state.rows = []
if "selected_template" not in st.session_state:
st.session_state.selected_template = "<select>"
if "markdown_log" not in st.session_state:
st.session_state.markdown_log = []
rows = st.session_state.rows
def download_function(project_name: str,
project_desc: str,
project_data: Dict[str, Any],
meeting: Dict[str, Any],
transcript) -> None:
files = export_meeting(project_name, project_desc, project_data["scientists"], meeting, transcript)
st.download_button("β¬οΈ DOCX", files["docx"],
file_name=f"{meeting['topic']}.docx",
mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
# project DB
def _load_projects() -> Dict[str, Any]:
if DB_FILE.exists():
return json.loads(DB_FILE.read_text())
return {}
def _save_projects(data: Dict[str, Any]):
DB_FILE.write_text(json.dumps(data, indent=2))
# scientist templates DB
def _load_templates() -> List[Dict[str, str]]:
if TEMPLATE_FILE.exists():
return json.loads(TEMPLATE_FILE.read_text())
# seed with three defaults on first run
defaults = [
{"title": "Immunologist", "expertise": "Immunopathology, antibodyβantigen interactions", "goal": "Guide immuneβtargeting strategies", "role": "Analyse epitope selection and immune response"},
{"title": "Machine Learning Expert", "expertise": "Deep learning, protein sequence modelling", "goal": "Develop predictive models for design", "role": "Build & chain ML models to rank candidates"},
{"title": "Computational Biologist", "expertise": "Protein folding simulation, molecular dynamics", "goal": "Validate structural stability", "role": "Simulate docking & refine structures"},
]
TEMPLATE_FILE.write_text(json.dumps(defaults, indent=2))
return defaults
def _save_templates(templates: List[Dict[str, str]]) -> None:
"""Persist scientist templates to disk."""
TEMPLATE_FILE.write_text(json.dumps(templates, indent=2))
def build_custom_thinktank(project_desc: str, scientists: List[Dict[str, str]]) -> ThinkTank:
lab = ThinkTank(project_desc)
lab.scientists.clear()
tools = [retrieve_documents]
for sd in scientists:
lab.scientists.append(
build_local_agent(
name=sd["title"],
description=f"Expertise: {sd['expertise']}. Goal: {sd['goal']}",
role=sd["role"],
memory=None,
storage=lab._storage,
tools=tools,
)
)
return lab
def write(name: str, content: str):
st.markdown(f"## π§βπ¬ {name} \n")
st.session_state.markdown_log.append(f"## π§βπ¬ {name} \n")
st.markdown(content)
st.session_state.markdown_log.append(content)
def default_serializer(obj):
"""Fallback serializer for JSON encoding."""
if hasattr(obj, "dict"): # Pydantic models
return obj.dict()
if hasattr(obj, "model_dump"): # Pydantic v2
return obj.model_dump()
if hasattr(obj, "__dict__"): # Generic Python objects
return obj.__dict__
return str(obj) # Final fallback
def save_response(response, filename_prefix="agent_response", project_name=None):
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
# Create project folder if project_name is provided
if project_name:
project_folder = clean_name(project_name)
os.makedirs(project_folder, exist_ok=True)
filename = os.path.join(project_folder, f"{filename_prefix}_{ts}.json")
else:
filename = f"{filename_prefix}_{ts}.json"
try:
with open(filename, "w", encoding="utf-8") as f:
# messages=response.messages
# collected = []
# for message in messages:
# if message.role == "tool" and message.tool_name == "retrieve_documents":
# retrieved_data_str = message.content
# retrieved_data = ast.literal_eval(retrieved_data_str)
# collected.append(retrieved_data)
json.dump(response, f, ensure_ascii=False, indent=2, default=default_serializer)
print(f"β
Saved response to {filename}")
except Exception as e:
print(f"β οΈ Failed to save response: {e}")
def run_thinktank_meeting(
project_name: str,
project_desc: str,
scientists: List[Dict[str, str]],
meeting_topic: str,
rounds: int,
projects_db: Dict[str, Any],
):
"""Execute a team meeting and write transcript + summary back to database."""
st.session_state.markdown_log = [] # reset log for new meeting
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = []
for scientist in st.session_state.rows:
print(f"Processing scientist: {scientist['title']}")
files = scientist.get('files', [])
print(f"Files for {scientist['title']}: {files}")
file_bytes_list = [(f.name, f.getvalue()) for f in files]
future = executor.submit(process_documents, file_bytes_list, clean_name(scientist['title']))
futures.append(future)
# Optionally wait for all to complete and handle exceptions
for future in concurrent.futures.as_completed(futures):
try:
future.result()
except Exception as e:
st.error(f"Error during ingestion: {e}")
lab = build_custom_thinktank(project_desc, scientists)
st.subheader(f"π§βπ¬ Team Meeting - {meeting_topic}")
st.session_state.markdown_log.append(f"# π§βπ¬ Team Meeting - {meeting_topic}")
def log(name: str, content: str):
write(name, content)
# PI opening
pi_open_response = lab.pi.run(
f"You are convening a team meeting. Agenda: {meeting_topic}. Share initial guidance to the scientists..",
stream=False,
)
pi_open = pi_open_response.content
log(lab.pi.name, pi_open)
save_response(pi_open_response, filename_prefix=f"{lab.pi.name}_opening", project_name=project_name)
# Discussion rounds
for r in range(1, rounds + 1):
st.markdown(f"### π Round {r}/{rounds}")
st.session_state.markdown_log.append(f"### π Round {r}/{rounds}")
for sci in lab.scientists:
tool_prompt = f"""
You are an expert in a team meeting. Your task is to contribute to the discussion based on your expertise and the context provided.
DO NOT summarize or paraphrase the context, but use it to inform your response.
Generate a new response every time.
You have access to the following tool:
1.Tool: `retrieve_documents`
- Purpose: Retrieve relevant document chunks from the knowledge database using natural language queries.
- Usage:
1. Analyze the current task or context and formulate meaningful queries.
2. Call: retrieve_documents(queries: List[str], collection_name: str) -> List[str]
3. Use collection_name = {clean_name(sci.name)}
Instructions:
- First, think about what information is needed to accomplish your task.
- Generate targeted, specific queries based on your expertise.
- Use `retrieve_documents` to fetch supporting content.
- Incorporate retrieved content directly into your reasoning or task output.
- **Do not output the summary or paraphrase the retrieved content β use it as-is.**
Your goal is to leverage the retrieved knowledge to solve the task accurately and completely.
"""
response = sci.run(
f"{tool_prompt}\n\nGive your contribution for round {r}:",
stream=False,
)
resp = response.content
lab._log("scientist", sci.name, resp)
log(sci.name, resp)
save_response(response, filename_prefix=f"{sci.name}_{r}_response", project_name=project_name)
crit = lab.critic.run(
f"Context so far:\n{lab._context()}\n\nCritique round {r}",
stream=False,
).content
save_response(crit, filename_prefix=f"{lab.critic.name}_critique", project_name=project_name)
lab._log("critic", lab.critic.name, crit)
log(lab.critic.name, crit)
synth = lab.pi.run(
f"Context so far:\n{lab._context()}\n\nSynthesise round {r} and pose followβups.",
stream=False,
).content
lab._log("pi", lab.pi.name, synth)
log(lab.pi.name + " (synthesis)", synth)
save_response(synth, filename_prefix=f"{lab.pi.name}_synthesis", project_name=project_name)
# Final summary
summary = lab.pi.run(
f"Context so far:\n{lab._context()}\n\nProvide the final detailed meeting summary and recommendations.",
stream=False,
).content
lab._log("pi", lab.pi.name, summary)
save_response(summary, filename_prefix=f"{lab.pi.name}_final_summary", project_name=project_name)
st.subheader("π Final Meeting Summary")
st.session_state.markdown_log.append("π Final Meeting Summary")
st.markdown(summary)
st.session_state.markdown_log.append(summary)
# Memory is now handled automatically by the agent's storage
# lab._memory.add_user_memory(memory=UserMemory(memory=summary), user_id=project_name)
# Save to DB
proj = projects_db.setdefault(project_name, {"description": project_desc, "scientists": scientists, "meetings": []})
for sci in scientists:
sci.pop("files", None) # remove files from scientists to avoid bloating the DB
proj["description"] = project_desc
proj["scientists"] = scientists
proj["meetings"].append({
"timestamp": int(time.time()),
"topic": meeting_topic,
"rounds": rounds,
"transcript": st.session_state.markdown_log,
"summary": summary,
})
_save_projects(projects_db)
st.success("Meeting complete and saved π")
return proj
st.set_page_config(page_title="Think Tank", layout="wide")
st.markdown(
"""
<style>
[data-testid="stSidebar"] {
width: 500px;
max-width: 1500px;
}
[data-testid="stSidebar"] + div .block-container {
padding-left: 300px;
}
</style>
""",
unsafe_allow_html=True,
)
st.markdown(
"<h1 style='text-align:center;'>π§ Think Tank</h1>",
unsafe_allow_html=True,
)
projects_db = _load_projects()
project_names = sorted(projects_db.keys())
# ββ Project selection / creation
st.sidebar.markdown(
f"""
<div style="display: flex; align-items: center;">
<img src="data:image/png;base64,{img_b64}" width="40" style="margin-right:10px;">
<span style="font-size: 14px;">Developed by TAU Group</span>
</div>
""",
unsafe_allow_html=True,
)
st.sidebar.markdown("---")
st.sidebar.header("Project Manager")
proj_choice = st.sidebar.selectbox("Select a Project", ["β New project"] + project_names)
if proj_choice == "β New project":
project_name = st.sidebar.text_input("New project name")
project_data = {"description": "", "scientists": [], "meetings": []}
if not project_name:
st.stop()
else:
project_name = proj_choice
project_data = projects_db[project_name]
# ββ Project description --
project_desc = st.sidebar.text_area("Project description", value=project_data.get("description", ""), height=120)
# ββ Scientist roster --
# Scientist templates loaded from disk
TEMPLATES: List[Dict[str, str]] = _load_templates()
st.sidebar.subheader("Experts Manager")
# Template management
with st.sidebar.expander("Manage Experts and templates", expanded=False):
# Select existing template to load
selected_tpl_title = st.selectbox("Load template to edit", ["<new>"] + [t["title"] for t in TEMPLATES], key="tpl_select")
# Populate fields depending on selection
if selected_tpl_title == "<new>":
tpl_data = {"title": "", "expertise": "", "goal": "", "role": ""}
else:
tpl_data = next(t for t in TEMPLATES if t["title"] == selected_tpl_title)
new_t = st.text_input("Title", value=tpl_data["title"], key="tpl_title")
col1, col2 = st.columns(2)
with col1:
new_e = st.text_area("Expertise", value=tpl_data["expertise"], height=70, key="tpl_exp")
with col2:
new_g = st.text_area("Goal", value=tpl_data["goal"], height=70, key="tpl_goal")
new_r = st.text_area("Role", value=tpl_data["role"], height=70, key="tpl_role")
# Save / update template
if st.button("Save template") and new_t.strip():
TEMPLATES = [t for t in TEMPLATES if t["title"] != new_t]
TEMPLATES.append({"title": new_t, "expertise": new_e, "goal": new_g, "role": new_r})
_save_templates(TEMPLATES)
(st.rerun if hasattr(st, 'rerun') else st.experimental_rerun)()
# delete templates
del_tpl = st.multiselect("Delete templates", [t["title"] for t in TEMPLATES])
if del_tpl and st.button("Remove selected templates"):
TEMPLATES = [t for t in TEMPLATES if t["title"] not in del_tpl]
_save_templates(TEMPLATES)
(st.rerun if hasattr(st, 'rerun') else st.experimental_rerun)()
# Load or initialise project rows project rows
if proj_choice == 'β New project':
if "rows" not in st.session_state:
st.session_state.rows = project_data.get("scientists", [])
else:
if len(st.session_state.rows) < 1 or 'files' not in st.session_state.rows[-1]:
st.session_state.rows = project_data.get("scientists", [])
def _add_template():
sel = st.session_state.tpl_selectbox
if sel != "<select>" and sel not in [r["title"] for r in st.session_state.rows]:
st.session_state.rows.append(next(t for t in TEMPLATES if t["title"] == sel))
_save_projects(projects_db)
st.session_state.tpl_selectbox = "<select>"
# Add from template --
st.selectbox(
"Add scientist from template",
["<select>"] + [t["title"] for t in TEMPLATES],
key="tpl_selectbox",
on_change=_add_template,
)
# Manual create scientist
if st.button("Add blank scientist"):
st.session_state.rows.append({"title": f"Scientist {len(st.session_state.rows)+1}", "expertise": "", "goal": "", "role": ""})
st.rerun()
# Editable table
st.session_state.rows = st.session_state.rows[:8] # limit to 8
display_rows = []
real_rows = st.session_state.rows
for r in real_rows:
r_view = r.copy()
if "files" in r_view:
# show just the count or filenames
val = r_view["files"]
r_view["files"] = len(val) if isinstance(val, (list, tuple)) else (val or 0)
display_rows.append(r_view)
scientist_table = st.data_editor(display_rows, num_rows="dynamic", use_container_width=True, key="scientist_table")
def _files_changed(i: int):
uploaded = st.session_state[f"files_{i}"] # new value after change
st.session_state.rows[i]["files"] = uploaded
print(f"Files for scientist {i}: {[f.name for f in uploaded]}")
for i, scientist in enumerate(st.session_state.rows):
c = st.container()
c.markdown(f"### Files for {scientist['title'] or f'Scientist {i+1}'}")
c.file_uploader(
"Choose files for the vector store",
accept_multiple_files=True,
key=f"files_{i}",
on_change=_files_changed,
args=(i,), # pass only the index
)
# ββ Meeting selection / creation -- / creation --
st.sidebar.subheader("Team Meeting")
meetings = project_data.get("meetings", [])
meeting_labels = [f"{i+1}. {m['topic']}" for i, m in enumerate(meetings)]
meeting_choice = st.sidebar.selectbox("Select a meeting", ["New meeting"] + meeting_labels)
if meeting_choice == "New meeting":
meeting_topic = st.sidebar.text_input("New meeting topic / title")
rounds = int(st.sidebar.number_input("Rounds", min_value=1, value=3, step=1))
run_btn = st.sidebar.button("Run Team Meeting")
if run_btn:
clean_scientists = [row for row in scientist_table if row["title"].strip()]
if not clean_scientists:
st.error("Provide at least one scientist")
st.stop()
proj = run_thinktank_meeting(project_name, project_desc, clean_scientists, meeting_topic, rounds, projects_db)
download_function(project_name, project_desc, project_data, proj["meetings"][-1], st.session_state.markdown_log)
else:
# Load existing meeting
sel_index = meeting_labels.index(meeting_choice)
meeting = meetings[sel_index]
st.session_state.markdown_log = []
st.markdown(f"## ποΈ Meeting Record - {meeting['topic']}")
st.session_state.markdown_log.append(f"## ποΈ Meeting Record - {meeting['topic']}")
for msg in meeting["transcript"]:
if type(msg) is dict:
name = msg.get("name", "Unknown")
content = msg.get("content", "")
write(name, content)
else:
st.markdown(msg)
st.session_state.markdown_log.append(msg)
st.markdown("### Summary")
st.session_state.markdown_log.append("### Summary")
st.markdown(meeting["summary"])
st.session_state.markdown_log.append(meeting["summary"])
download_function(project_name, project_desc, project_data, meeting, meeting['transcript'])