llm4smeft 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- llm4smeft-0.1.0/PKG-INFO +50 -0
- llm4smeft-0.1.0/README.md +29 -0
- llm4smeft-0.1.0/llm4smeft/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/benchmarks/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/benchmarks/cases.py +25 -0
- llm4smeft-0.1.0/llm4smeft/chat/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/chat/session.py +256 -0
- llm4smeft-0.1.0/llm4smeft/chat/store.py +93 -0
- llm4smeft-0.1.0/llm4smeft/data/EWPO.json +1011 -0
- llm4smeft-0.1.0/llm4smeft/data/EW_VV.json +1329 -0
- llm4smeft-0.1.0/llm4smeft/data/EW_all.json +1332 -0
- llm4smeft-0.1.0/llm4smeft/data/LEP_WW.json +414 -0
- llm4smeft-0.1.0/llm4smeft/data/LHC_VV.json +874 -0
- llm4smeft-0.1.0/llm4smeft/data/VV_all.json +902 -0
- llm4smeft-0.1.0/llm4smeft/data/W_helicity.json +117 -0
- llm4smeft-0.1.0/llm4smeft/data/global.json +5706 -0
- llm4smeft-0.1.0/llm4smeft/data/higgs_EW.json +3072 -0
- llm4smeft-0.1.0/llm4smeft/data/higgs_SS.json +718 -0
- llm4smeft-0.1.0/llm4smeft/data/higgs_STXS.json +1191 -0
- llm4smeft-0.1.0/llm4smeft/data/higgs_VV.json +2491 -0
- llm4smeft-0.1.0/llm4smeft/data/higgs_all.json +1571 -0
- llm4smeft-0.1.0/llm4smeft/data/higgs_hh.json +84 -0
- llm4smeft-0.1.0/llm4smeft/data/single_top.json +344 -0
- llm4smeft-0.1.0/llm4smeft/data/singletop_all.json +692 -0
- llm4smeft-0.1.0/llm4smeft/data/tW.json +141 -0
- llm4smeft-0.1.0/llm4smeft/data/tZ.json +239 -0
- llm4smeft-0.1.0/llm4smeft/data/top_4heavy.json +262 -0
- llm4smeft-0.1.0/llm4smeft/data/top_EW.json +4240 -0
- llm4smeft-0.1.0/llm4smeft/data/top_all.json +2844 -0
- llm4smeft-0.1.0/llm4smeft/data/top_higgs.json +4331 -0
- llm4smeft-0.1.0/llm4smeft/data/ttV.json +989 -0
- llm4smeft-0.1.0/llm4smeft/data/ttW.json +139 -0
- llm4smeft-0.1.0/llm4smeft/data/ttZ.json +699 -0
- llm4smeft-0.1.0/llm4smeft/data/tt_incl.json +869 -0
- llm4smeft-0.1.0/llm4smeft/data/tta.json +104 -0
- llm4smeft-0.1.0/llm4smeft/data/ttbb.json +151 -0
- llm4smeft-0.1.0/llm4smeft/data/tttt.json +164 -0
- llm4smeft-0.1.0/llm4smeft/evaluator/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/evaluator/backend.py +56 -0
- llm4smeft-0.1.0/llm4smeft/evaluator/results.py +16 -0
- llm4smeft-0.1.0/llm4smeft/experiments/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/experiments/run_interactive.py +776 -0
- llm4smeft-0.1.0/llm4smeft/hypotheses/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/hypotheses/schema.py +20 -0
- llm4smeft-0.1.0/llm4smeft/prompts/base_prompt.txt +59 -0
- llm4smeft-0.1.0/llm4smeft/proposers/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/proposers/base.py +25 -0
- llm4smeft-0.1.0/llm4smeft/proposers/llm.py +516 -0
- llm4smeft-0.1.0/llm4smeft/proposers/parsing.py +259 -0
- llm4smeft-0.1.0/llm4smeft/proposers/prompting.py +82 -0
- llm4smeft-0.1.0/llm4smeft/summaries/__init__.py +0 -0
- llm4smeft-0.1.0/llm4smeft/summaries/retrieval.py +450 -0
- llm4smeft-0.1.0/llm4smeft/summaries/schema.py +15 -0
- llm4smeft-0.1.0/llm4smeft.egg-info/PKG-INFO +50 -0
- llm4smeft-0.1.0/llm4smeft.egg-info/SOURCES.txt +60 -0
- llm4smeft-0.1.0/llm4smeft.egg-info/dependency_links.txt +1 -0
- llm4smeft-0.1.0/llm4smeft.egg-info/entry_points.txt +2 -0
- llm4smeft-0.1.0/llm4smeft.egg-info/requires.txt +10 -0
- llm4smeft-0.1.0/llm4smeft.egg-info/top_level.txt +1 -0
- llm4smeft-0.1.0/pyproject.toml +41 -0
- llm4smeft-0.1.0/setup.cfg +4 -0
- llm4smeft-0.1.0/tests/test_pipeline.py +180 -0
llm4smeft-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: llm4smeft
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: LLM-driven interactive hypothesis generation for SMEFT global fits
|
|
5
|
+
Author: HEP
|
|
6
|
+
License: MIT
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: Operating System :: OS Independent
|
|
9
|
+
Requires-Python: >=3.10
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
Requires-Dist: numpy<2
|
|
12
|
+
Requires-Dist: pandas
|
|
13
|
+
Requires-Dist: huggingface-hub<1.0.0,>=0.33.4
|
|
14
|
+
Requires-Dist: transformers<4.52.0,>=4.45.0
|
|
15
|
+
Requires-Dist: accelerate<1.0.0,>=0.26.0
|
|
16
|
+
Requires-Dist: langchain>=0.3.0
|
|
17
|
+
Requires-Dist: langchain-huggingface<1.3.0,>=0.1.0
|
|
18
|
+
Requires-Dist: pydantic>=2.0.0
|
|
19
|
+
Requires-Dist: torch
|
|
20
|
+
Requires-Dist: python-dotenv
|
|
21
|
+
|
|
22
|
+
# LLM4SMEFT
|
|
23
|
+
|
|
24
|
+
LLM-driven interactive hypothesis generation for SMEFT global fits. It works as a prior hyptheses inference to SMEFiT package.
|
|
25
|
+
|
|
26
|
+
## Install
|
|
27
|
+
|
|
28
|
+
```bash
|
|
29
|
+
pip install llm4smeft
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
## Run
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
llm4smeft
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
This launches the interactive session (equivalent to the old
|
|
40
|
+
`./run.sh` / `python src/experiments/run_interactive.py`).
|
|
41
|
+
|
|
42
|
+
## Data
|
|
43
|
+
|
|
44
|
+
The SMEFiT summary records used by the retrieval module ship inside the
|
|
45
|
+
package (`llm4smeft/data/*.json`) and are located relative to the
|
|
46
|
+
installed package files, not the current working directory — so
|
|
47
|
+
`llm4smeft` works from any directory after `pip install`. Session
|
|
48
|
+
records accepted during a run are written back into that same
|
|
49
|
+
installed `data/` folder unless `LLM4SMEFT_DATA_DIR` is set to another
|
|
50
|
+
path.
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
# LLM4SMEFT
|
|
2
|
+
|
|
3
|
+
LLM-driven interactive hypothesis generation for SMEFT global fits. It works as a prior hyptheses inference to SMEFiT package.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install llm4smeft
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
## Run
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
llm4smeft
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
This launches the interactive session (equivalent to the old
|
|
19
|
+
`./run.sh` / `python src/experiments/run_interactive.py`).
|
|
20
|
+
|
|
21
|
+
## Data
|
|
22
|
+
|
|
23
|
+
The SMEFiT summary records used by the retrieval module ship inside the
|
|
24
|
+
package (`llm4smeft/data/*.json`) and are located relative to the
|
|
25
|
+
installed package files, not the current working directory — so
|
|
26
|
+
`llm4smeft` works from any directory after `pip install`. Session
|
|
27
|
+
records accepted during a run are written back into that same
|
|
28
|
+
installed `data/` folder unless `LLM4SMEFT_DATA_DIR` is set to another
|
|
29
|
+
path.
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""Data structures for benchmark cases."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from dataclasses import dataclass, field
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
ALLOWED_CASE_TYPES = {"synthetic", "uv", "null"}
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass(slots=True)
|
|
12
|
+
class BenchmarkCase:
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
case_id: str
|
|
16
|
+
case_type: str
|
|
17
|
+
true_support: list[str] | None = None
|
|
18
|
+
true_coeffs: dict[str, float] | None = None
|
|
19
|
+
metadata: dict[str, object] = field(default_factory=dict)
|
|
20
|
+
|
|
21
|
+
def __post_init__(self) -> None:
|
|
22
|
+
"""Validate that the benchmark category is supported."""
|
|
23
|
+
if self.case_type not in ALLOWED_CASE_TYPES:
|
|
24
|
+
allowed = ", ".join(sorted(ALLOWED_CASE_TYPES))
|
|
25
|
+
raise ValueError(f"case_type must be one of: {allowed}")
|
|
File without changes
|
|
@@ -0,0 +1,256 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import uuid
|
|
4
|
+
from dataclasses import dataclass, field
|
|
5
|
+
from datetime import datetime, timezone
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
@dataclass
|
|
10
|
+
class TurnRecord:
|
|
11
|
+
"""All data produced in one propose → feedback round."""
|
|
12
|
+
turn: int
|
|
13
|
+
role: str # "user" | "assistant"
|
|
14
|
+
content: str # raw text sent / received
|
|
15
|
+
subtype: str | None = None # "summary" | "feedback" | None
|
|
16
|
+
reasoning: str | None = None # CoT reasoning block (assistant only)
|
|
17
|
+
hypotheses: list[dict] = field(default_factory=list)
|
|
18
|
+
evaluations: list[dict] = field(default_factory=list)
|
|
19
|
+
timestamp: str = field(default_factory=lambda: _now())
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
# ---------------------------------------------
|
|
23
|
+
# ChatSession
|
|
24
|
+
# ---------------------------------------------
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class ChatSession:
|
|
28
|
+
"""Manages the message list for one multi-turn SMEFT interaction.
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
def __init__(
|
|
32
|
+
self,
|
|
33
|
+
case_id: str,
|
|
34
|
+
model: str,
|
|
35
|
+
allowed_operators: list[str] | None = None,
|
|
36
|
+
max_support: int = 3,
|
|
37
|
+
session_id: str | None = None,
|
|
38
|
+
) -> None:
|
|
39
|
+
self.session_id: str = session_id or str(uuid.uuid4())[:8]
|
|
40
|
+
self.case_id: str = case_id
|
|
41
|
+
self.model: str = model
|
|
42
|
+
self.allowed_operators: list[str] = allowed_operators or []
|
|
43
|
+
self.max_support: int = max_support
|
|
44
|
+
self.started_at: str = _now()
|
|
45
|
+
self.finished_at: str | None = None
|
|
46
|
+
self.accepted_turn: int | None = None
|
|
47
|
+
|
|
48
|
+
self._messages: list[dict[str, str]] = []
|
|
49
|
+
self._records: list[TurnRecord] = []
|
|
50
|
+
self.current_turn: int = 0
|
|
51
|
+
|
|
52
|
+
# ---------------------------------------------
|
|
53
|
+
# Building the conversation
|
|
54
|
+
# ---------------------------------------------
|
|
55
|
+
|
|
56
|
+
def set_system_message(self, content: str) -> None:
|
|
57
|
+
"""Set (or replace) the system message. Called once at session start."""
|
|
58
|
+
self._messages = [m for m in self._messages if m["role"] != "system"]
|
|
59
|
+
self._messages.insert(0, {"role": "system", "content": content})
|
|
60
|
+
|
|
61
|
+
def add_summary_message(self, content: str) -> None:
|
|
62
|
+
"""Add the initial physics summary as the first user message."""
|
|
63
|
+
self.current_turn = 1
|
|
64
|
+
self._messages.append({"role": "user", "content": content})
|
|
65
|
+
self._records.append(TurnRecord(
|
|
66
|
+
turn=self.current_turn,
|
|
67
|
+
role="user",
|
|
68
|
+
content=content,
|
|
69
|
+
subtype="summary",
|
|
70
|
+
))
|
|
71
|
+
|
|
72
|
+
def add_feedback(self, feedback: str) -> None:
|
|
73
|
+
"""Add a user feedback message, wrapped with full context for the LLM.
|
|
74
|
+
|
|
75
|
+
"""
|
|
76
|
+
self.current_turn += 1
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
enriched = self._build_feedback_message(feedback)
|
|
80
|
+
|
|
81
|
+
self._messages.append({"role": "user", "content": enriched})
|
|
82
|
+
# Store the raw user text in records (clean for human reading)
|
|
83
|
+
self._records.append(TurnRecord(
|
|
84
|
+
turn=self.current_turn,
|
|
85
|
+
role="user",
|
|
86
|
+
content=feedback,
|
|
87
|
+
subtype="feedback",
|
|
88
|
+
))
|
|
89
|
+
|
|
90
|
+
def _build_feedback_message(self, raw_feedback: str) -> str:
|
|
91
|
+
|
|
92
|
+
prev_hyps = self.last_hypotheses()
|
|
93
|
+
if prev_hyps:
|
|
94
|
+
lines = []
|
|
95
|
+
for i, h in enumerate(prev_hyps, 1):
|
|
96
|
+
ops = ", ".join(h.get("operators", []))
|
|
97
|
+
tag = h.get("theory_tag") or ""
|
|
98
|
+
rat = h.get("rationale") or ""
|
|
99
|
+
lines.append(f" {i}. operators=[{ops}] tag={tag}")
|
|
100
|
+
if rat:
|
|
101
|
+
lines.append(f' "{rat}"')
|
|
102
|
+
hyp_block = "\n".join(lines)
|
|
103
|
+
else:
|
|
104
|
+
hyp_block = " (none)"
|
|
105
|
+
|
|
106
|
+
ops_str = ", ".join(self.allowed_operators) if self.allowed_operators else "(not set)"
|
|
107
|
+
|
|
108
|
+
return (
|
|
109
|
+
f"[REVISION REQUEST — Turn {self.current_turn}]\n\n"
|
|
110
|
+
f"Your previous hypotheses:\n{hyp_block}\n\n"
|
|
111
|
+
f"User feedback: {raw_feedback}\n\n"
|
|
112
|
+
f"Please revise your hypotheses based on this feedback.\n"
|
|
113
|
+
f"Constraints remain unchanged:\n"
|
|
114
|
+
f" - Allowed operators: {ops_str}\n"
|
|
115
|
+
f" - At most {self.max_support} operators per hypothesis\n"
|
|
116
|
+
f" - Return the <reasoning>...</reasoning>"
|
|
117
|
+
f"<hypotheses>...</hypotheses> format"
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
def add_assistant_turn(
|
|
121
|
+
self,
|
|
122
|
+
raw_response: str,
|
|
123
|
+
reasoning: str | None,
|
|
124
|
+
hypotheses: list[Any],
|
|
125
|
+
evaluations: list[Any] | None = None,
|
|
126
|
+
) -> None:
|
|
127
|
+
"""Record the LLM's response for the current turn."""
|
|
128
|
+
self._messages.append({"role": "assistant", "content": raw_response or ""})
|
|
129
|
+
self._records.append(TurnRecord(
|
|
130
|
+
turn=self.current_turn,
|
|
131
|
+
role="assistant",
|
|
132
|
+
content=raw_response,
|
|
133
|
+
reasoning=reasoning,
|
|
134
|
+
hypotheses=[_hyp_to_dict(h) for h in hypotheses],
|
|
135
|
+
evaluations=[_eval_to_dict(e) for e in (evaluations or [])],
|
|
136
|
+
))
|
|
137
|
+
|
|
138
|
+
def close(self, accepted_turn: int | None = None) -> None:
|
|
139
|
+
"""Mark the session as finished."""
|
|
140
|
+
self.finished_at = _now()
|
|
141
|
+
self.accepted_turn = accepted_turn
|
|
142
|
+
|
|
143
|
+
# ---------------------------------------------
|
|
144
|
+
# Accessors
|
|
145
|
+
# ---------------------------------------------
|
|
146
|
+
|
|
147
|
+
def messages_for_llm(self, flatten: bool = False) -> list[dict[str, str]]:
|
|
148
|
+
|
|
149
|
+
if not flatten:
|
|
150
|
+
return list(self._messages)
|
|
151
|
+
|
|
152
|
+
system = [m for m in self._messages if m["role"] == "system"]
|
|
153
|
+
return system + [{"role": "user", "content": self.combined_user_message()}]
|
|
154
|
+
|
|
155
|
+
def combined_user_message(self) -> str:
|
|
156
|
+
|
|
157
|
+
summary_msg = next(
|
|
158
|
+
(r.content for r in self._records if r.subtype == "summary"), "")
|
|
159
|
+
|
|
160
|
+
transcript: list[str] = []
|
|
161
|
+
for rec in self._records:
|
|
162
|
+
if rec.subtype == "summary":
|
|
163
|
+
continue
|
|
164
|
+
if rec.role == "assistant":
|
|
165
|
+
if rec.hypotheses:
|
|
166
|
+
lines = []
|
|
167
|
+
for i, h in enumerate(rec.hypotheses, 1):
|
|
168
|
+
ops = ", ".join(h.get("operators", []))
|
|
169
|
+
tag = h.get("theory_tag") or ""
|
|
170
|
+
lines.append(f" {i}) operators ({ops}), tag: {tag}")
|
|
171
|
+
transcript.append(
|
|
172
|
+
f"[Turn {rec.turn} — your hypotheses]\n" + "\n".join(lines))
|
|
173
|
+
else:
|
|
174
|
+
transcript.append(
|
|
175
|
+
f"[Turn {rec.turn} — your hypotheses]\n (none parsed)")
|
|
176
|
+
elif rec.subtype == "feedback":
|
|
177
|
+
transcript.append(
|
|
178
|
+
f"[Turn {rec.turn} — user feedback]\n {rec.content}")
|
|
179
|
+
|
|
180
|
+
ops_str = ", ".join(self.allowed_operators) \
|
|
181
|
+
if self.allowed_operators else "(not set)"
|
|
182
|
+
|
|
183
|
+
parts = [summary_msg]
|
|
184
|
+
if transcript:
|
|
185
|
+
parts.append(
|
|
186
|
+
"--- Conversation so far (compact notation, for reference "
|
|
187
|
+
"ONLY — do NOT reply in this notation) ---\n"
|
|
188
|
+
+ "\n\n".join(transcript))
|
|
189
|
+
parts.append(
|
|
190
|
+
"[Current request]\n"
|
|
191
|
+
"Revise your hypotheses now, taking EVERY piece of user feedback "
|
|
192
|
+
"above into account (not only the most recent one).\n"
|
|
193
|
+
"Constraints remain unchanged:\n"
|
|
194
|
+
f" - Allowed operators: {ops_str}\n"
|
|
195
|
+
f" - At most {self.max_support} operators per hypothesis\n"
|
|
196
|
+
"Respond in EXACTLY the output format required by the system "
|
|
197
|
+
"instructions — the four sentinels, each exactly once, "
|
|
198
|
+
"nothing before or after:\n"
|
|
199
|
+
"<<<REASONING>>>\n"
|
|
200
|
+
"(your analysis — one step per line, each starting with '- ')\n"
|
|
201
|
+
"<<<END_REASONING>>>\n"
|
|
202
|
+
"<<<HYPOTHESES>>>\n"
|
|
203
|
+
"(the JSON object only)\n"
|
|
204
|
+
"<<<END_HYPOTHESES>>>"
|
|
205
|
+
)
|
|
206
|
+
return "\n\n".join(p for p in parts if p)
|
|
207
|
+
|
|
208
|
+
def last_hypotheses(self) -> list[dict]:
|
|
209
|
+
"""Return the hypotheses from the most recent assistant turn."""
|
|
210
|
+
for record in reversed(self._records):
|
|
211
|
+
if record.role == "assistant":
|
|
212
|
+
return record.hypotheses
|
|
213
|
+
return []
|
|
214
|
+
|
|
215
|
+
def n_turns(self) -> int:
|
|
216
|
+
return self.current_turn
|
|
217
|
+
|
|
218
|
+
# ------------------------------------------------------
|
|
219
|
+
# Serialisation
|
|
220
|
+
# ------------------------------------------------------
|
|
221
|
+
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
# ------------------------------------------------------
|
|
225
|
+
# Helpers
|
|
226
|
+
# ------------------------------------------------------
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
def _now() -> str:
|
|
230
|
+
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
def _hyp_to_dict(h: Any) -> dict:
|
|
234
|
+
if isinstance(h, dict):
|
|
235
|
+
return h
|
|
236
|
+
return {
|
|
237
|
+
"operators": getattr(h, "operators", []),
|
|
238
|
+
"relations": getattr(h, "relations", None),
|
|
239
|
+
"theory_tag": getattr(h, "theory_tag", None),
|
|
240
|
+
"rationale": getattr(h, "rationale", None),
|
|
241
|
+
}
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
def _eval_to_dict(e: Any) -> dict:
|
|
245
|
+
if isinstance(e, dict):
|
|
246
|
+
return e
|
|
247
|
+
return {
|
|
248
|
+
"chi2": getattr(e, "chi2", None),
|
|
249
|
+
"aic": getattr(e, "aic", None),
|
|
250
|
+
"bic": getattr(e, "bic", None),
|
|
251
|
+
"n_params": getattr(e, "n_params", None),
|
|
252
|
+
"metadata": getattr(e, "metadata", {}),
|
|
253
|
+
}
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
__all__ = ["ChatSession"]
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
"""ChatHistoryStore — saves completed ChatSession objects to disk as JSON files.
|
|
2
|
+
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
from __future__ import annotations
|
|
6
|
+
|
|
7
|
+
import json
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
from .session import ChatSession
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class ChatHistoryStore:
|
|
14
|
+
"""Save and load chat sessions to/from a directory of JSON files.
|
|
15
|
+
|
|
16
|
+
Parameters
|
|
17
|
+
----------
|
|
18
|
+
directory:
|
|
19
|
+
Path to the storage directory. Created automatically if absent.
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
def __init__(self, directory: str | Path = "chat_history") -> None:
|
|
23
|
+
self.directory = Path(directory)
|
|
24
|
+
self.directory.mkdir(parents=True, exist_ok=True)
|
|
25
|
+
|
|
26
|
+
# ------------------------------------------------------
|
|
27
|
+
# Save
|
|
28
|
+
# ------------------------------------------------------
|
|
29
|
+
|
|
30
|
+
def save(self, session: ChatSession) -> Path:
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
ts = session.started_at.replace(":", "-")
|
|
34
|
+
filename = f"{_safe(session.case_id)}_{session.session_id}_{ts}.json"
|
|
35
|
+
path = self.directory / filename
|
|
36
|
+
|
|
37
|
+
document = {
|
|
38
|
+
"meta": {
|
|
39
|
+
"session_id": session.session_id,
|
|
40
|
+
"case_id": session.case_id,
|
|
41
|
+
"model": session.model,
|
|
42
|
+
"started_at": session.started_at,
|
|
43
|
+
"finished_at": session.finished_at,
|
|
44
|
+
"total_turns": session.n_turns(),
|
|
45
|
+
"accepted_turn": session.accepted_turn,
|
|
46
|
+
},
|
|
47
|
+
"messages": [
|
|
48
|
+
{
|
|
49
|
+
"turn": rec.turn,
|
|
50
|
+
"role": rec.role,
|
|
51
|
+
"subtype": rec.subtype,
|
|
52
|
+
"content": rec.content,
|
|
53
|
+
"timestamp": rec.timestamp,
|
|
54
|
+
"reasoning": rec.reasoning,
|
|
55
|
+
"hypotheses": rec.hypotheses,
|
|
56
|
+
"evaluations": rec.evaluations,
|
|
57
|
+
}
|
|
58
|
+
for rec in session._records
|
|
59
|
+
],
|
|
60
|
+
"summary": {
|
|
61
|
+
"total_turns": session.n_turns(),
|
|
62
|
+
"accepted_turn": session.accepted_turn,
|
|
63
|
+
"final_hypotheses": session.last_hypotheses(),
|
|
64
|
+
},
|
|
65
|
+
}
|
|
66
|
+
|
|
67
|
+
path.write_text(json.dumps(document, indent=2), encoding="utf-8")
|
|
68
|
+
print(f"Session saved → {path}")
|
|
69
|
+
return path
|
|
70
|
+
|
|
71
|
+
# ------------------------------------------------------
|
|
72
|
+
# Load
|
|
73
|
+
# ------------------------------------------------------
|
|
74
|
+
|
|
75
|
+
def load(self, path: str | Path) -> dict:
|
|
76
|
+
"""Load a single session file. Returns the parsed JSON document."""
|
|
77
|
+
return json.loads(Path(path).read_text(encoding="utf-8"))
|
|
78
|
+
|
|
79
|
+
def list_sessions(self) -> list[Path]:
|
|
80
|
+
"""Return sorted list of all session file paths."""
|
|
81
|
+
return sorted(self.directory.glob("*.json"))
|
|
82
|
+
|
|
83
|
+
def load_all(self) -> list[dict]:
|
|
84
|
+
"""Load and return all session files as a list of dicts."""
|
|
85
|
+
return [self.load(p) for p in self.list_sessions()]
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def _safe(text: str) -> str:
|
|
90
|
+
return "".join(c if c.isalnum() or c in "-_" else "_" for c in text)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
__all__ = ["ChatHistoryStore"]
|