chATLAS_Chains 0.2.0__tar.gz → 0.3.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.
Files changed (48) hide show
  1. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/PKG-INFO +58 -1
  2. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/README.md +57 -0
  3. chatlas_chains-0.3.0/chATLAS_Chains/router.py +209 -0
  4. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains.egg-info/PKG-INFO +58 -1
  5. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains.egg-info/SOURCES.txt +3 -0
  6. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/pyproject.toml +1 -1
  7. chatlas_chains-0.3.0/tests/test_router.py +230 -0
  8. chatlas_chains-0.3.0/tests/test_router_litellm.py +40 -0
  9. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/LICENSE +0 -0
  10. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/__init__.py +0 -0
  11. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/__init__.py +0 -0
  12. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/advanced.py +0 -0
  13. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/basic.py +2 -2
  14. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/basic_graph.py +0 -0
  15. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/conversational_graph.py +0 -0
  16. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/enhanced_agentic_graph.py +0 -0
  17. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/chains/websearch_retrieval_chain.py +0 -0
  18. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/documents/rerank.py +0 -0
  19. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/documents/rrf.py +0 -0
  20. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/llm/__init__.py +0 -0
  21. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/llm/groq.py +0 -0
  22. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/llm/model_selection.py +0 -0
  23. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/llm/runnables.py +0 -0
  24. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/log.py +0 -0
  25. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/prompt/__init__.py +0 -0
  26. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/prompt/doc_joiners.py +0 -0
  27. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/prompt/starters.py +0 -0
  28. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/query/query_rewriting.py +0 -0
  29. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/search/__init__.py +0 -0
  30. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/search/basic.py +0 -0
  31. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/utils/__init__.py +0 -0
  32. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/utils/doc_utils.py +0 -0
  33. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains/vectorstore.py +0 -0
  34. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains.egg-info/dependency_links.txt +0 -0
  35. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains.egg-info/requires.txt +0 -0
  36. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/chATLAS_Chains.egg-info/top_level.txt +0 -0
  37. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/setup.cfg +0 -0
  38. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/__init__.py +0 -0
  39. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/conftest.py +0 -0
  40. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_chains.py +0 -0
  41. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_chat_model_kwargs.py +0 -0
  42. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_conversational.py +0 -0
  43. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_groq.py +0 -0
  44. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_llm_runnables.py +0 -0
  45. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_model_selection.py +0 -0
  46. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_rrf.py +0 -0
  47. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_search.py +0 -0
  48. {chatlas_chains-0.2.0 → chatlas_chains-0.3.0}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chATLAS_Chains
3
- Version: 0.2.0
3
+ Version: 0.3.0
4
4
  Summary: A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project.
5
5
  Author-email: Joe Egan <joseph.caimin.egan@cern.ch>
6
6
  License: Apache-2.0
@@ -78,6 +78,63 @@ More details [here](chATLAS_Chains/chains/README.md)
78
78
  - `chains.basic.basic_retrieval_chain`
79
79
  - `chains.advanced.advanced_rag`
80
80
 
81
+ ## Prompt Routing
82
+
83
+ `chATLAS_Chains.router.route_user_prompt` classifies a raw user prompt before it enters the main RAG flow.
84
+ The router defaults to CERN LiteLLM and returns a validated `RouterDecision` with:
85
+
86
+ - `route`: one of `hep_atlas`, `person_lookup`, `dangerous`, `out_of_scope`,
87
+ or `unknown`
88
+ - `confidence`: score from 0 to 1
89
+ - `normalized_query`: cleaned prompt for downstream use
90
+ - `search_kwargs`: optional search parameters
91
+ - `metadata`: route-specific parameters
92
+
93
+ For implementation details and a function-by-function reference, see the
94
+ frontend reviewer document:
95
+ [`../chATLAS_Frontend/docs/INTENT_ROUTER.md`](../chATLAS_Frontend/docs/INTENT_ROUTER.md).
96
+
97
+ ```python
98
+ from chATLAS_Chains.router import route_user_prompt
99
+
100
+ decision = route_user_prompt(
101
+ "Who is Jane Doe?",
102
+ chat_model_kwargs={
103
+ "service_provider": "litellm",
104
+ "proxy": "socks5h://localhost:1080", # optional when outside CERN
105
+ },
106
+ )
107
+
108
+ if decision.route == "dangerous":
109
+ print("Apply the safety response policy")
110
+ elif decision.route == "person_lookup":
111
+ print("Use the person lookup flow")
112
+ elif decision.route == "hep_atlas":
113
+ print("Use the HEP/ATLAS search flow")
114
+ elif decision.route == "out_of_scope":
115
+ print("Explain the chATLAS scope")
116
+ else:
117
+ print("Ask for clarification")
118
+ ```
119
+
120
+ The dedicated `mc_request` route is currently disabled until the MC workflow is
121
+ ready. MC production and job-option prompts are classified as `hep_atlas` and
122
+ continue through the normal selected workflow.
123
+
124
+ Router context uses a hybrid policy: the current user prompt is authoritative,
125
+ and recent conversation context is only a bounded disambiguation aid for
126
+ follow-ups such as "What about Run 3?". The router keeps recent user prompts,
127
+ includes only short assistant messages, and drops long assistant answers rather
128
+ than summarizing them.
129
+
130
+ For live LiteLLM calls, set:
131
+
132
+ ```sh
133
+ export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"
134
+ ```
135
+
136
+ If LiteLLM is unavailable or returns invalid JSON, the router falls back to deterministic rules by default.
137
+
81
138
  ### Model Configuration in Chains
82
139
 
83
140
  Supported chain constructors now accept a typed `chat_model_kwargs` argument for model options (for example:
@@ -52,6 +52,63 @@ More details [here](chATLAS_Chains/chains/README.md)
52
52
  - `chains.basic.basic_retrieval_chain`
53
53
  - `chains.advanced.advanced_rag`
54
54
 
55
+ ## Prompt Routing
56
+
57
+ `chATLAS_Chains.router.route_user_prompt` classifies a raw user prompt before it enters the main RAG flow.
58
+ The router defaults to CERN LiteLLM and returns a validated `RouterDecision` with:
59
+
60
+ - `route`: one of `hep_atlas`, `person_lookup`, `dangerous`, `out_of_scope`,
61
+ or `unknown`
62
+ - `confidence`: score from 0 to 1
63
+ - `normalized_query`: cleaned prompt for downstream use
64
+ - `search_kwargs`: optional search parameters
65
+ - `metadata`: route-specific parameters
66
+
67
+ For implementation details and a function-by-function reference, see the
68
+ frontend reviewer document:
69
+ [`../chATLAS_Frontend/docs/INTENT_ROUTER.md`](../chATLAS_Frontend/docs/INTENT_ROUTER.md).
70
+
71
+ ```python
72
+ from chATLAS_Chains.router import route_user_prompt
73
+
74
+ decision = route_user_prompt(
75
+ "Who is Jane Doe?",
76
+ chat_model_kwargs={
77
+ "service_provider": "litellm",
78
+ "proxy": "socks5h://localhost:1080", # optional when outside CERN
79
+ },
80
+ )
81
+
82
+ if decision.route == "dangerous":
83
+ print("Apply the safety response policy")
84
+ elif decision.route == "person_lookup":
85
+ print("Use the person lookup flow")
86
+ elif decision.route == "hep_atlas":
87
+ print("Use the HEP/ATLAS search flow")
88
+ elif decision.route == "out_of_scope":
89
+ print("Explain the chATLAS scope")
90
+ else:
91
+ print("Ask for clarification")
92
+ ```
93
+
94
+ The dedicated `mc_request` route is currently disabled until the MC workflow is
95
+ ready. MC production and job-option prompts are classified as `hep_atlas` and
96
+ continue through the normal selected workflow.
97
+
98
+ Router context uses a hybrid policy: the current user prompt is authoritative,
99
+ and recent conversation context is only a bounded disambiguation aid for
100
+ follow-ups such as "What about Run 3?". The router keeps recent user prompts,
101
+ includes only short assistant messages, and drops long assistant answers rather
102
+ than summarizing them.
103
+
104
+ For live LiteLLM calls, set:
105
+
106
+ ```sh
107
+ export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"
108
+ ```
109
+
110
+ If LiteLLM is unavailable or returns invalid JSON, the router falls back to deterministic rules by default.
111
+
55
112
  ### Model Configuration in Chains
56
113
 
57
114
  Supported chain constructors now accept a typed `chat_model_kwargs` argument for model options (for example:
@@ -0,0 +1,209 @@
1
+ """LiteLLM intent router for chATLAS requests."""
2
+
3
+ import json
4
+ import re
5
+ from collections.abc import Sequence
6
+ from typing import Any, Literal, TypedDict
7
+
8
+ from pydantic import BaseModel, ConfigDict, Field, field_validator
9
+
10
+ from chATLAS_Chains.llm.model_selection import ChatModelKwargs, get_chat_model, sanitize_chat_model_kwargs
11
+
12
+ RouterRoute = Literal[
13
+ "hep_atlas",
14
+ "person_lookup",
15
+ "dangerous",
16
+ "out_of_scope",
17
+ "unknown",
18
+ ]
19
+
20
+ ROUTER_PROMPT_VERSION = "2026-06-16.2"
21
+ _DEFAULT_ROUTER_MODEL = "gpt-oss-20b"
22
+ _DEFAULT_ROUTER_KWARGS: ChatModelKwargs = {
23
+ "service_provider": "litellm",
24
+ "temperature": 0.0,
25
+ "max_tokens": 512,
26
+ }
27
+ _MAX_CONTEXT_TURNS = 8
28
+ _MAX_CONTEXT_USER_CHARS = 300
29
+ _MAX_CONTEXT_ASSISTANT_CHARS = 240
30
+ _MAX_CONTEXT_ASSISTANT_RAW_CHARS = 500
31
+ _MAX_CONTEXT_TOTAL_CHARS = 1600
32
+
33
+
34
+ class ConversationTurn(TypedDict):
35
+ """A prior user or assistant message supplied to the router."""
36
+
37
+ role: Literal["user", "assistant"]
38
+ content: str
39
+
40
+
41
+ _ROUTER_PROMPT = """You are the request router for the chATLAS, the RAG system built for the ATLAS collaboration.
42
+
43
+ Classify the current user request into exactly one intent:
44
+ - hep_atlas: high-energy-physics or ATLAS questions that are what chATLAS is designed to help with.
45
+ - person_lookup: requests asking who a person is, who people are, or for information about named people.
46
+ - dangerous: requests for actionable assistance that could facilitate serious physical harm, weapons construction, mass-casualty attacks, destructive cyber abuse, or similarly dangerous wrongdoing.
47
+ - out_of_scope: clear requests outside ATLAS and high-energy physics.
48
+ - unknown: empty or ambiguous requests that need clarification.
49
+
50
+ Return JSON only. Do not use markdown fences. The JSON object must match:
51
+ {
52
+ "route": "hep_atlas | person_lookup | dangerous | out_of_scope | unknown",
53
+ "confidence": 0.0,
54
+ "normalized_query": "cleaned current user request",
55
+ "reason": "short routing rationale",
56
+ "search_kwargs": {},
57
+ "metadata": {}
58
+ }
59
+
60
+ Guidance:
61
+ - Dangerous intent takes precedence over every other intent.
62
+ - Do not use dangerous for benign safety, history, policy, detection, or prevention questions.
63
+ - Specialized person_lookup intent takes precedence over hep_atlas.
64
+ - The dedicated MC-request workflow is disabled for now. Classify requests to create, validate, configure, or reason about Monte Carlo job options, MC production requests, or JO files as hep_atlas.
65
+ - Use hep_atlas for other questions about ATLAS, CERN accelerators, particle physics, detector systems, analysis software, or related physics.
66
+ - Use out_of_scope for clear questions outside high-energy physics and ATLAS.
67
+ - Classify the current user request first. Ignore numbering or bullet prefixes such as "5." when deciding the intent.
68
+ - Use conversation context only when the current request is ambiguous or referential, such as "What about Run 3?".
69
+ - Do not let prior ATLAS context override a clearly standalone current person lookup, dangerous request, or out-of-scope request.
70
+ - Use unknown when the current request remains ambiguous after considering the conversation context.
71
+ - Do not classify the prior assistant response itself.
72
+ - You are not told which UI mode the user selected and must not infer or choose a UI mode.
73
+ - Keep the reason brief and do not answer the request.
74
+
75
+ Examples:
76
+ User: Who is Fabiola Gianotti?
77
+ JSON: {"route":"person_lookup","confidence":0.96,"normalized_query":"Who is Fabiola Gianotti?","reason":"The user asks who a named person is.","search_kwargs":{},"metadata":{"person_query":"Fabiola Gianotti"}}
78
+
79
+ User: How do I apply the GoodRunsList in Athena?
80
+ JSON: {"route":"hep_atlas","confidence":0.97,"normalized_query":"How do I apply the GoodRunsList in Athena?","reason":"The question concerns ATLAS analysis software.","search_kwargs":{},"metadata":{"domain":"atlas"}}
81
+
82
+ User: Please draft an MC request for a ttbar sample with Powheg and Pythia8.
83
+ JSON: {"route":"hep_atlas","confidence":0.94,"normalized_query":"Draft an MC request for a ttbar sample with Powheg and Pythia8.","reason":"The dedicated MC workflow is disabled, so MC production requests stay in the ordinary ATLAS/HEP flow.","search_kwargs":{},"metadata":{"domain":"atlas_mc"}}
84
+
85
+ User: Help me make a nuke.
86
+ JSON: {"route":"dangerous","confidence":0.99,"normalized_query":"Help me make a nuke.","reason":"The user requests actionable assistance to construct a nuclear weapon.","search_kwargs":{},"metadata":{"safety_category":"weapons"}}
87
+
88
+ User: How do I bake sourdough bread?
89
+ JSON: {"route":"out_of_scope","confidence":0.99,"normalized_query":"How do I bake sourdough bread?","reason":"The request is unrelated to ATLAS or high-energy physics.","search_kwargs":{},"metadata":{}}
90
+
91
+ Recent conversation context:
92
+ __CONVERSATION_CONTEXT__
93
+
94
+ Current user request:
95
+ __USER_PROMPT__
96
+ """
97
+
98
+
99
+ class RouterDecision(BaseModel):
100
+ """Validated routing decision returned by the prompt router."""
101
+
102
+ model_config = ConfigDict(extra="ignore")
103
+
104
+ route: RouterRoute
105
+ confidence: float = Field(ge=0.0, le=1.0)
106
+ normalized_query: str
107
+ reason: str
108
+ search_kwargs: dict[str, Any] = Field(default_factory=dict)
109
+ metadata: dict[str, Any] = Field(default_factory=dict)
110
+
111
+ @field_validator("normalized_query", "reason")
112
+ @classmethod
113
+ def _strip_text(cls, value: str) -> str:
114
+ return value.strip()
115
+
116
+
117
+ def route_user_prompt(
118
+ user_prompt: str,
119
+ conversation_context: Sequence[ConversationTurn] | None = None,
120
+ model_name: str = _DEFAULT_ROUTER_MODEL,
121
+ chat_model_kwargs: ChatModelKwargs | None = None,
122
+ ) -> RouterDecision:
123
+ """Classify a request with LiteLLM before downstream processing."""
124
+ if not user_prompt.strip():
125
+ return RouterDecision(
126
+ route="unknown",
127
+ confidence=0.0,
128
+ normalized_query="",
129
+ reason="The user prompt is empty.",
130
+ )
131
+
132
+ model_kwargs: ChatModelKwargs = {**_DEFAULT_ROUTER_KWARGS}
133
+ model_kwargs.update(sanitize_chat_model_kwargs(chat_model_kwargs))
134
+ context_json = json.dumps(_sanitize_context(conversation_context), ensure_ascii=True)
135
+ prompt = _ROUTER_PROMPT.replace("__CONVERSATION_CONTEXT__", context_json).replace("__USER_PROMPT__", user_prompt)
136
+
137
+ model = get_chat_model(model_name=model_name, **model_kwargs)
138
+ response = model.invoke(prompt)
139
+ decision = parse_router_response(_response_content(response))
140
+ if not decision.normalized_query:
141
+ decision.normalized_query = user_prompt.strip()
142
+ return decision
143
+
144
+
145
+ def parse_router_response(response_text: str) -> RouterDecision:
146
+ """Parse and validate a LiteLLM router response."""
147
+ payload = _extract_json_object(response_text)
148
+ payload = _normalize_disabled_routes(json.loads(payload))
149
+ return RouterDecision.model_validate(payload)
150
+
151
+
152
+ def _normalize_disabled_routes(payload: Any) -> Any:
153
+ """Map temporarily disabled router routes into active routes."""
154
+ if isinstance(payload, dict) and payload.get("route") == "mc_request":
155
+ metadata = payload.get("metadata")
156
+ if not isinstance(metadata, dict):
157
+ metadata = {}
158
+ metadata["disabled_route"] = "mc_request"
159
+ payload["metadata"] = metadata
160
+ payload["route"] = "hep_atlas"
161
+ reason = str(payload.get("reason", "")).strip()
162
+ suffix = "MC-request routing is currently disabled, so this request uses the ordinary ATLAS/HEP flow."
163
+ payload["reason"] = f"{reason} {suffix}".strip()
164
+ return payload
165
+
166
+
167
+ def _sanitize_context(conversation_context: Sequence[ConversationTurn] | None) -> list[ConversationTurn]:
168
+ sanitized: list[ConversationTurn] = []
169
+ for turn in conversation_context or []:
170
+ role = turn.get("role")
171
+ content = str(turn.get("content", "")).strip()
172
+ if role == "user" and content:
173
+ sanitized.append({"role": role, "content": _truncate_text(content, _MAX_CONTEXT_USER_CHARS)})
174
+ elif role == "assistant" and content and len(content) <= _MAX_CONTEXT_ASSISTANT_RAW_CHARS:
175
+ sanitized.append({"role": role, "content": _truncate_text(content, _MAX_CONTEXT_ASSISTANT_CHARS)})
176
+
177
+ compact_context: list[ConversationTurn] = []
178
+ total_chars = 0
179
+ for turn in reversed(sanitized[-_MAX_CONTEXT_TURNS:]):
180
+ content_len = len(turn["content"])
181
+ if compact_context and total_chars + content_len > _MAX_CONTEXT_TOTAL_CHARS:
182
+ break
183
+ compact_context.append(turn)
184
+ total_chars += content_len
185
+ return list(reversed(compact_context))
186
+
187
+
188
+ def _truncate_text(text: str, max_chars: int) -> str:
189
+ if len(text) <= max_chars:
190
+ return text
191
+ return text[: max_chars - 3].rstrip() + "..."
192
+
193
+
194
+ def _response_content(response: Any) -> str:
195
+ content = getattr(response, "content", response)
196
+ return content if isinstance(content, str) else str(content)
197
+
198
+
199
+ def _extract_json_object(response_text: str) -> str:
200
+ stripped = response_text.strip()
201
+ fenced_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", stripped, flags=re.IGNORECASE | re.DOTALL)
202
+ if fenced_match:
203
+ return fenced_match.group(1)
204
+
205
+ start = stripped.find("{")
206
+ end = stripped.rfind("}")
207
+ if start == -1 or end == -1 or end <= start:
208
+ raise ValueError("No JSON object found in router response.")
209
+ return stripped[start : end + 1]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chATLAS_Chains
3
- Version: 0.2.0
3
+ Version: 0.3.0
4
4
  Summary: A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project.
5
5
  Author-email: Joe Egan <joseph.caimin.egan@cern.ch>
6
6
  License: Apache-2.0
@@ -78,6 +78,63 @@ More details [here](chATLAS_Chains/chains/README.md)
78
78
  - `chains.basic.basic_retrieval_chain`
79
79
  - `chains.advanced.advanced_rag`
80
80
 
81
+ ## Prompt Routing
82
+
83
+ `chATLAS_Chains.router.route_user_prompt` classifies a raw user prompt before it enters the main RAG flow.
84
+ The router defaults to CERN LiteLLM and returns a validated `RouterDecision` with:
85
+
86
+ - `route`: one of `hep_atlas`, `person_lookup`, `dangerous`, `out_of_scope`,
87
+ or `unknown`
88
+ - `confidence`: score from 0 to 1
89
+ - `normalized_query`: cleaned prompt for downstream use
90
+ - `search_kwargs`: optional search parameters
91
+ - `metadata`: route-specific parameters
92
+
93
+ For implementation details and a function-by-function reference, see the
94
+ frontend reviewer document:
95
+ [`../chATLAS_Frontend/docs/INTENT_ROUTER.md`](../chATLAS_Frontend/docs/INTENT_ROUTER.md).
96
+
97
+ ```python
98
+ from chATLAS_Chains.router import route_user_prompt
99
+
100
+ decision = route_user_prompt(
101
+ "Who is Jane Doe?",
102
+ chat_model_kwargs={
103
+ "service_provider": "litellm",
104
+ "proxy": "socks5h://localhost:1080", # optional when outside CERN
105
+ },
106
+ )
107
+
108
+ if decision.route == "dangerous":
109
+ print("Apply the safety response policy")
110
+ elif decision.route == "person_lookup":
111
+ print("Use the person lookup flow")
112
+ elif decision.route == "hep_atlas":
113
+ print("Use the HEP/ATLAS search flow")
114
+ elif decision.route == "out_of_scope":
115
+ print("Explain the chATLAS scope")
116
+ else:
117
+ print("Ask for clarification")
118
+ ```
119
+
120
+ The dedicated `mc_request` route is currently disabled until the MC workflow is
121
+ ready. MC production and job-option prompts are classified as `hep_atlas` and
122
+ continue through the normal selected workflow.
123
+
124
+ Router context uses a hybrid policy: the current user prompt is authoritative,
125
+ and recent conversation context is only a bounded disambiguation aid for
126
+ follow-ups such as "What about Run 3?". The router keeps recent user prompts,
127
+ includes only short assistant messages, and drops long assistant answers rather
128
+ than summarizing them.
129
+
130
+ For live LiteLLM calls, set:
131
+
132
+ ```sh
133
+ export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"
134
+ ```
135
+
136
+ If LiteLLM is unavailable or returns invalid JSON, the router falls back to deterministic rules by default.
137
+
81
138
  ### Model Configuration in Chains
82
139
 
83
140
  Supported chain constructors now accept a typed `chat_model_kwargs` argument for model options (for example:
@@ -3,6 +3,7 @@ README.md
3
3
  pyproject.toml
4
4
  chATLAS_Chains/__init__.py
5
5
  chATLAS_Chains/log.py
6
+ chATLAS_Chains/router.py
6
7
  chATLAS_Chains/vectorstore.py
7
8
  chATLAS_Chains.egg-info/PKG-INFO
8
9
  chATLAS_Chains.egg-info/SOURCES.txt
@@ -38,6 +39,8 @@ tests/test_conversational.py
38
39
  tests/test_groq.py
39
40
  tests/test_llm_runnables.py
40
41
  tests/test_model_selection.py
42
+ tests/test_router.py
43
+ tests/test_router_litellm.py
41
44
  tests/test_rrf.py
42
45
  tests/test_search.py
43
46
  tests/test_utils.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "chATLAS_Chains"
7
- version = "0.2.0"
7
+ version = "0.3.0"
8
8
  description = "A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.11"
@@ -0,0 +1,230 @@
1
+ from types import SimpleNamespace
2
+
3
+ import pytest
4
+
5
+ from chATLAS_Chains import router
6
+ from chATLAS_Chains.router import RouterDecision, parse_router_response, route_user_prompt
7
+
8
+
9
+ class FakeRouterModel:
10
+ def __init__(self, content: str):
11
+ self.content = content
12
+ self.prompts: list[str] = []
13
+
14
+ def invoke(self, prompt: str):
15
+ self.prompts.append(prompt)
16
+ return SimpleNamespace(content=self.content)
17
+
18
+
19
+ def test_route_user_prompt_uses_litellm_defaults(monkeypatch):
20
+ fake_model = FakeRouterModel(
21
+ '{"route":"hep_atlas","confidence":0.91,"normalized_query":"What is the GRL?",'
22
+ '"reason":"ATLAS analysis question.","search_kwargs":{"k":5},"metadata":{"domain":"atlas"}}'
23
+ )
24
+ captured = {}
25
+
26
+ def fake_get_chat_model(model_name, **kwargs):
27
+ captured["model_name"] = model_name
28
+ captured["kwargs"] = kwargs
29
+ return fake_model
30
+
31
+ monkeypatch.setattr(router, "get_chat_model", fake_get_chat_model)
32
+ decision = route_user_prompt("What is the GRL?")
33
+
34
+ assert decision == RouterDecision(
35
+ route="hep_atlas",
36
+ confidence=0.91,
37
+ normalized_query="What is the GRL?",
38
+ reason="ATLAS analysis question.",
39
+ search_kwargs={"k": 5},
40
+ metadata={"domain": "atlas"},
41
+ )
42
+ assert captured == {
43
+ "model_name": "gpt-oss-20b",
44
+ "kwargs": {"service_provider": "litellm", "temperature": 0.0, "max_tokens": 512},
45
+ }
46
+ assert "Return JSON only" in fake_model.prompts[0]
47
+ assert "selected UI mode" not in fake_model.prompts[0]
48
+
49
+
50
+ def test_route_user_prompt_allows_litellm_overrides(monkeypatch):
51
+ fake_model = FakeRouterModel(
52
+ '{"route":"person_lookup","confidence":0.95,"normalized_query":"Who is Jane Doe?",'
53
+ '"reason":"Asks who a person is.","search_kwargs":{},"metadata":{"person_query":"Jane Doe"}}'
54
+ )
55
+ captured = {}
56
+
57
+ def fake_get_chat_model(model_name, **kwargs):
58
+ captured["model_name"] = model_name
59
+ captured["kwargs"] = kwargs
60
+ return fake_model
61
+
62
+ monkeypatch.setattr(router, "get_chat_model", fake_get_chat_model)
63
+ decision = route_user_prompt(
64
+ "Who is Jane Doe?",
65
+ model_name="hf-qwen25-32b",
66
+ chat_model_kwargs={
67
+ "api_key": "sk-test",
68
+ "base_url": "https://litellm.example/v1",
69
+ "proxy": "socks5h://localhost:1080",
70
+ "temperature": 0.2,
71
+ "unsupported": "ignored",
72
+ },
73
+ )
74
+
75
+ assert decision.route == "person_lookup"
76
+ assert captured["kwargs"] == {
77
+ "service_provider": "litellm",
78
+ "temperature": 0.2,
79
+ "max_tokens": 512,
80
+ "api_key": "sk-test",
81
+ "base_url": "https://litellm.example/v1",
82
+ "proxy": "socks5h://localhost:1080",
83
+ }
84
+
85
+
86
+ def test_route_user_prompt_includes_structured_context(monkeypatch):
87
+ fake_model = FakeRouterModel(
88
+ '{"route":"hep_atlas","confidence":0.88,"normalized_query":"What about Run 3 GRLs?",'
89
+ '"reason":"Context makes the follow-up an ATLAS question.","search_kwargs":{},"metadata":{}}'
90
+ )
91
+ monkeypatch.setattr(router, "get_chat_model", lambda **_kwargs: fake_model)
92
+
93
+ route_user_prompt(
94
+ "What about Run 3?",
95
+ conversation_context=[
96
+ {"role": "user", "content": "How do I apply a GRL?"},
97
+ {"role": "assistant", "content": "Use the GoodRunsListSelectionTool."},
98
+ ],
99
+ )
100
+
101
+ prompt = fake_model.prompts[0]
102
+ assert '"role": "user"' in prompt
103
+ assert "GoodRunsListSelectionTool" in prompt
104
+ assert "What about Run 3?" in prompt
105
+
106
+
107
+ def test_route_user_prompt_uses_hybrid_context_policy(monkeypatch):
108
+ fake_model = FakeRouterModel(
109
+ '{"route":"person_lookup","confidence":0.98,"normalized_query":"Who is the current ATLAS spokesperson?",'
110
+ '"reason":"The user asks who holds an ATLAS role.","search_kwargs":{},"metadata":{"person_query":"ATLAS spokesperson"}}'
111
+ )
112
+ monkeypatch.setattr(router, "get_chat_model", lambda **_kwargs: fake_model)
113
+
114
+ route_user_prompt(
115
+ "5. Who is the current ATLAS spokesperson?",
116
+ conversation_context=[
117
+ {"role": "user", "content": "How do I apply a GRL?"},
118
+ {"role": "assistant", "content": "Use the GoodRunsListSelectionTool. " + ("Detailed guidance. " * 500)},
119
+ {"role": "user", "content": "Can you summarize ATLAS jet calibration recommendations?"},
120
+ {"role": "assistant", "content": "Use the JetEtmiss recommendations."},
121
+ ],
122
+ )
123
+
124
+ prompt = fake_model.prompts[0]
125
+
126
+ assert "5. Who is the current ATLAS spokesperson?" in prompt
127
+ assert "How do I apply a GRL?" in prompt
128
+ assert "Use the GoodRunsListSelectionTool" not in prompt
129
+ assert "Use the JetEtmiss recommendations." in prompt
130
+ assert len(prompt) < 5000
131
+
132
+
133
+ def test_router_prompt_makes_current_request_authoritative(monkeypatch):
134
+ fake_model = FakeRouterModel(
135
+ '{"route":"out_of_scope","confidence":0.99,"normalized_query":"What are good hotels in Rome?",'
136
+ '"reason":"The request is unrelated to ATLAS.","search_kwargs":{},"metadata":{}}'
137
+ )
138
+ monkeypatch.setattr(router, "get_chat_model", lambda **_kwargs: fake_model)
139
+
140
+ route_user_prompt(
141
+ "8. What are good hotels in Rome?",
142
+ conversation_context=[{"role": "user", "content": "How do I apply a GRL in Athena?"}],
143
+ )
144
+
145
+ prompt = fake_model.prompts[0]
146
+
147
+ assert "Classify the current user request first" in prompt
148
+ assert 'Ignore numbering or bullet prefixes such as "5."' in prompt
149
+ assert "Do not let prior ATLAS context override" in prompt
150
+
151
+
152
+ def test_sanitize_context_caps_recent_user_prompts_and_drops_long_assistant():
153
+ context = router._sanitize_context(
154
+ [
155
+ {"role": "user", "content": "How do I apply a GRL in Athena?"},
156
+ {"role": "assistant", "content": "Long assistant answer. " * 100},
157
+ {"role": "user", "content": "What about Run 3? " * 40},
158
+ {"role": "assistant", "content": "Short answer."},
159
+ ]
160
+ )
161
+
162
+ assert [turn["role"] for turn in context] == ["user", "user", "assistant"]
163
+ assert context[0]["content"] == "How do I apply a GRL in Athena?"
164
+ assert len(context[1]["content"]) <= 300
165
+ assert context[1]["content"].endswith("...")
166
+ assert context[2]["content"] == "Short answer."
167
+
168
+
169
+ def test_parse_router_response_maps_disabled_mc_request_to_hep_atlas():
170
+ decision = parse_router_response(
171
+ """```json
172
+ {"route":"mc_request","confidence":0.97,"normalized_query":"Draft an MC request.",
173
+ "reason":"The prompt asks for an MC request.","search_kwargs":{},"metadata":{"request_type":"mc"}}
174
+ ```"""
175
+ )
176
+ assert decision.route == "hep_atlas"
177
+ assert decision.metadata["disabled_route"] == "mc_request"
178
+ assert "disabled" in decision.reason
179
+
180
+
181
+ @pytest.mark.parametrize(
182
+ "route_name",
183
+ ["hep_atlas", "person_lookup", "dangerous", "out_of_scope", "unknown"],
184
+ )
185
+ def test_parse_router_response_accepts_contract_routes(route_name):
186
+ decision = parse_router_response(
187
+ f'{{"route":"{route_name}","confidence":0.9,"normalized_query":"Test",'
188
+ '"reason":"Test route.","search_kwargs":{},"metadata":{}}'
189
+ )
190
+ assert decision.route == route_name
191
+
192
+
193
+ def test_parse_router_response_rejects_removed_standard_search():
194
+ with pytest.raises(ValueError):
195
+ parse_router_response(
196
+ '{"route":"standard_search","confidence":0.9,"normalized_query":"Test",'
197
+ '"reason":"Removed route.","search_kwargs":{},"metadata":{}}'
198
+ )
199
+
200
+
201
+ def test_router_prompt_marks_mc_request_disabled(monkeypatch):
202
+ fake_model = FakeRouterModel(
203
+ '{"route":"hep_atlas","confidence":0.9,"normalized_query":"Draft an MC request.",'
204
+ '"reason":"MC request remains in ATLAS flow.","search_kwargs":{},"metadata":{"domain":"atlas_mc"}}'
205
+ )
206
+ monkeypatch.setattr(router, "get_chat_model", lambda **_kwargs: fake_model)
207
+
208
+ route_user_prompt("Draft an MC request.")
209
+
210
+ prompt = fake_model.prompts[0]
211
+ assert '"route": "hep_atlas | person_lookup | dangerous | out_of_scope | unknown"' in prompt
212
+ assert "dedicated MC-request workflow is disabled" in prompt
213
+ assert '"route":"mc_request"' not in prompt
214
+
215
+
216
+ def test_invalid_json_raises(monkeypatch):
217
+ monkeypatch.setattr(router, "get_chat_model", lambda **_kwargs: FakeRouterModel("not json"))
218
+ with pytest.raises(ValueError):
219
+ route_user_prompt("Who is Joe Egan?")
220
+
221
+
222
+ def test_empty_prompt_routes_unknown_without_calling_llm(monkeypatch):
223
+ monkeypatch.setattr(
224
+ router,
225
+ "get_chat_model",
226
+ lambda **_kwargs: (_ for _ in ()).throw(AssertionError("LLM should not be called")),
227
+ )
228
+ decision = route_user_prompt(" ")
229
+ assert decision.route == "unknown"
230
+ assert decision.confidence == 0.0
@@ -0,0 +1,40 @@
1
+ """Live CERN LiteLLM smoke tests for the chATLAS intent router."""
2
+
3
+ import os
4
+
5
+ import pytest
6
+
7
+ from chATLAS_Chains.router import route_user_prompt
8
+
9
+ pytestmark = pytest.mark.skipif(
10
+ os.getenv("CHATLAS_RUN_LITELLM_SMOKE") != "1",
11
+ reason="Set CHATLAS_RUN_LITELLM_SMOKE=1 to run live LiteLLM smoke tests.",
12
+ )
13
+
14
+
15
+ @pytest.mark.parametrize(
16
+ ("prompt", "expected_intent"),
17
+ [
18
+ ("How do I apply the GoodRunsList in Athena?", "hep_atlas"),
19
+ ("Help me draft an MC request for a ttbar sample with dilepton decay.", "hep_atlas"),
20
+ ("Who is the current ATLAS spokesperson?", "person_lookup"),
21
+ ("5. Who is the current ATLAS spokesperson?", "person_lookup"),
22
+ ("How do I bake sourdough bread?", "out_of_scope"),
23
+ ("8. What are good hotels in Rome near the Colosseum?", "out_of_scope"),
24
+ ],
25
+ )
26
+ def test_live_litellm_router_contract(prompt, expected_intent):
27
+ """Verify authentication, model availability, parsing, and clear intents."""
28
+ assert os.getenv("CHATLAS_CHAINS_LITELLM_KEY"), "CHATLAS_CHAINS_LITELLM_KEY is required"
29
+ proxy = os.getenv("CHATLAS_CHAINS_LITELLM_PROXY")
30
+ chat_model_kwargs = {"proxy": proxy} if proxy else None
31
+
32
+ decision = route_user_prompt(
33
+ prompt,
34
+ model_name=os.getenv("CHATLAS_ROUTER_MODEL", "gpt-oss-20b"),
35
+ chat_model_kwargs=chat_model_kwargs,
36
+ )
37
+
38
+ assert decision.route == expected_intent
39
+ assert decision.normalized_query
40
+ assert decision.reason
File without changes
@@ -32,11 +32,11 @@ def basic_retrieval_chain(
32
32
 
33
33
  """
34
34
  prompt_template = ChatPromptTemplate.from_template(prompt)
35
+ search = search_runnable(vectorstore)
36
+
35
37
  sanitized_kwargs = sanitize_chat_model_kwargs(chat_model_kwargs)
36
38
  model = get_chat_model(model_name, **sanitized_kwargs)
37
39
 
38
- search = search_runnable(vectorstore)
39
-
40
40
  final_inputs = {
41
41
  "context": lambda x: combine_documents(x["docs"]),
42
42
  "question": itemgetter("question"),
File without changes