chATLAS_Chains 0.1.3__py3-none-any.whl
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.
- benchmark/basic.py +34 -0
- benchmark/conversational.py +0 -0
- chATLAS_Chains/__init__.py +0 -0
- chATLAS_Chains/chains/__init__.py +0 -0
- chATLAS_Chains/chains/basic.py +42 -0
- chATLAS_Chains/chains/basic_graph.py +163 -0
- chATLAS_Chains/chains/conversational.py +0 -0
- chATLAS_Chains/chains/websearch_retrieval_chain.py +447 -0
- chATLAS_Chains/llm/__init__.py +0 -0
- chATLAS_Chains/llm/model_selection.py +47 -0
- chATLAS_Chains/log.py +20 -0
- chATLAS_Chains/prompt/__init__.py +0 -0
- chATLAS_Chains/prompt/doc_joiners.py +5 -0
- chATLAS_Chains/prompt/starters.py +151 -0
- chATLAS_Chains/search/__init__.py +0 -0
- chATLAS_Chains/search/basic.py +45 -0
- chATLAS_Chains/utils/__init__.py +0 -0
- chATLAS_Chains/utils/doc_utils.py +28 -0
- chATLAS_Chains/vectorstore.py +94 -0
- chatlas_chains-0.1.3.dist-info/METADATA +113 -0
- chatlas_chains-0.1.3.dist-info/RECORD +30 -0
- chatlas_chains-0.1.3.dist-info/WHEEL +5 -0
- chatlas_chains-0.1.3.dist-info/licenses/LICENSE +201 -0
- chatlas_chains-0.1.3.dist-info/top_level.txt +3 -0
- tests/__init__.py +0 -0
- tests/conftest.py +274 -0
- tests/test_chains.py +54 -0
- tests/test_llm.py +81 -0
- tests/test_search.py +48 -0
- tests/test_utils.py +27 -0
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
import importlib
|
|
2
|
+
import os
|
|
3
|
+
|
|
4
|
+
SUPPORTED_CHAT_MODELS = {
|
|
5
|
+
"gpt-3.5-turbo": {"library": "langchain_openai", "module": "ChatOpenAI"},
|
|
6
|
+
"gpt-4": {"library": "langchain_openai", "module": "ChatOpenAI"},
|
|
7
|
+
"gpt-4o": {"library": "langchain_openai", "module": "ChatOpenAI"},
|
|
8
|
+
"gpt-4o-mini": {"library": "langchain_openai", "module": "ChatOpenAI"},
|
|
9
|
+
}
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def get_chat_model(model_name):
|
|
13
|
+
"""
|
|
14
|
+
Initialize chat model with the provided model name (if supported)
|
|
15
|
+
|
|
16
|
+
This function dynamically imports the appropriate library and model class based on
|
|
17
|
+
the given `model_name`, initializes it using the `CHATLAS_OPENAI_KEY` from the environment,
|
|
18
|
+
and returns the model instance. If the model name is not supported, it defaults to `"gpt-4o-mini"`.
|
|
19
|
+
|
|
20
|
+
:param model_name: The name of the model to load (e.g., "gpt-4", "gpt-3.5-turbo").
|
|
21
|
+
:type model_name: str
|
|
22
|
+
|
|
23
|
+
:raises ValueError: If the environment variable `CHATLAS_OPENAI_KEY` is not set.
|
|
24
|
+
|
|
25
|
+
:return: An instance of the specified chat model.
|
|
26
|
+
:rtype: BaseLanguageModel
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
if model_name not in SUPPORTED_CHAT_MODELS:
|
|
30
|
+
model_name = "gpt-4o-mini"
|
|
31
|
+
|
|
32
|
+
api_key = os.getenv("CHATLAS_OPENAI_KEY")
|
|
33
|
+
if not api_key:
|
|
34
|
+
raise ValueError("CHATLAS_OPENAI_KEY not set in environment")
|
|
35
|
+
api_key = api_key.strip()
|
|
36
|
+
|
|
37
|
+
model_config = SUPPORTED_CHAT_MODELS[model_name]
|
|
38
|
+
library = model_config["library"]
|
|
39
|
+
module = model_config["module"]
|
|
40
|
+
|
|
41
|
+
library = importlib.import_module(library)
|
|
42
|
+
model_class = getattr(library, module)
|
|
43
|
+
|
|
44
|
+
# Dynamically create the model instance
|
|
45
|
+
model = model_class(model_name=model_name, openai_api_key=api_key)
|
|
46
|
+
|
|
47
|
+
return model
|
chATLAS_Chains/log.py
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
|
|
3
|
+
logger = logging.getLogger("chATLAS_Chains")
|
|
4
|
+
logger.setLevel(logging.DEBUG)
|
|
5
|
+
|
|
6
|
+
# Console Handler
|
|
7
|
+
console_handler = logging.StreamHandler()
|
|
8
|
+
console_handler.setLevel(logging.DEBUG)
|
|
9
|
+
console_formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
|
10
|
+
console_handler.setFormatter(console_formatter)
|
|
11
|
+
|
|
12
|
+
# File Handler
|
|
13
|
+
file_handler = logging.FileHandler("chATLAS_Chains.log")
|
|
14
|
+
file_handler.setLevel(logging.INFO)
|
|
15
|
+
file_formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
|
|
16
|
+
file_handler.setFormatter(file_formatter)
|
|
17
|
+
|
|
18
|
+
# Add handlers
|
|
19
|
+
logger.addHandler(console_handler)
|
|
20
|
+
logger.addHandler(file_handler)
|
|
File without changes
|
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
CHAT_PROMPT_TEMPLATE = """
|
|
2
|
+
You are an elite-level physicist specializing in high-energy physics, specifically within the ATLAS collaboration at CERN.
|
|
3
|
+
Your role is to serve as a highly reliable expert assistant, providing precise, accurate, and contextually relevant answers to professional physicists' inquiries.
|
|
4
|
+
|
|
5
|
+
### Instructions:
|
|
6
|
+
- **Use Only Provided Context:** You must strictly base your answers on the supplied context. Do not speculate, hallucinate, or fabricate information.
|
|
7
|
+
- **Answer with Clarity and Precision:** Provide clear, concise, and technically accurate responses. Use professional and scientific language appropriate for CERN physicists.
|
|
8
|
+
- **Format and Cite Properly:**
|
|
9
|
+
- Use Markdown for structured formatting (e.g., `**bold**` for emphasis, lists for steps, and code blocks for technical content).
|
|
10
|
+
- Explicitly cite all sources used by including the exact extract(s) you referenced, along with the document name and date in IEEE style. Use consistent citation formatting.
|
|
11
|
+
- **Show Reasoning and Steps:** When applicable, break down complex answers into clear, logical steps. Include derivations, explanations of methods, or detailed interpretations of the physics concepts involved.
|
|
12
|
+
- **Handle Ambiguity:** If the question or context is ambiguous, specify the missing details and suggest clarifying questions.
|
|
13
|
+
- **Use Structured Outputs When Relevant:** When the response involves multiple parameters, datasets, or step-by-step instructions, format the output using tables, lists, or clearly labeled sections for readability.
|
|
14
|
+
- **Prioritize Caution and Verifiability:** When the context is limited, explicitly state the lack of sufficient information rather than making unsupported inferences.
|
|
15
|
+
- **Context Prioritization:** When multiple sources provide relevant information, prioritize: (1) official ATLAS documentation, (2) peer-reviewed publications, (3) technical notes, (4) meeting minutes. When sources conflict, explicitly note the discrepancy and indicate which source appears most authoritative or recent.
|
|
16
|
+
- **Temporal Awareness:** Be attentive to publication dates in the provided context. Clearly indicate when information might be outdated and prioritize the most recent sources when appropriate. Flag any potential inconsistencies between older and newer documentation.
|
|
17
|
+
- **Cross-Reference Information:** When multiple sources address the same topic, synthesize the information by cross-referencing details from all relevant sources. Explicitly note where sources complement or contradict each other.
|
|
18
|
+
- **Technical Terminology:** Maintain consistency with ATLAS-specific conventions for detector components, physics objects, software tools, and analysis techniques. Use standard HEP units and follow ATLAS notation for particles, processes, and statistical methods.
|
|
19
|
+
- **Multi-Stage Physics Reasoning:** For complex questions, employ systematic physics reasoning by: (1) identifying relevant principles, (2) establishing assumptions and boundary conditions, (3) developing a logical sequence of calculations or deductions, and (4) relating the conclusion back to the original question in the context of ATLAS physics.
|
|
20
|
+
- **Express Uncertainty Appropriately:** Clearly distinguish between: established facts, widely accepted interpretations, active areas of research, and speculative explanations. Use precise language that reflects the confidence level of each statement (e.g., "definitively shown" vs. "suggests" vs. "remains unclear").
|
|
21
|
+
- **Context Limitation Awareness:** Explicitly identify knowledge gaps in the provided sources that affect your answer. Distinguish between: (1) information not included in the context, (2) genuinely open questions in the field, and (3) information that may exist but isn't provided.
|
|
22
|
+
- **Represent Data Appropriately:** When discussing experimental results, clearly describe relevant plots, significance values, and uncertainties. For detector components or event topologies, use clear descriptive language that could be translated to visual representation.
|
|
23
|
+
- **Include Relevant Metadata:** When referencing specific ATLAS runs, datasets, or analyses, include key metadata such as: run numbers, √s values, integrated luminosity, data-taking periods, and software/calibration versions when available in the context.
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
### Context (from ATLAS documentation and wiki pages):
|
|
28
|
+
{context}
|
|
29
|
+
|
|
30
|
+
### Question:
|
|
31
|
+
{question}
|
|
32
|
+
|
|
33
|
+
### Response:
|
|
34
|
+
"""
|
|
35
|
+
|
|
36
|
+
WEB_SEARCH_CHAT_PROMPT_TEMPLATE = """
|
|
37
|
+
You are an elite-level physicist specializing in high-energy physics, specifically within the ATLAS collaboration at CERN.
|
|
38
|
+
Your role is to serve as a highly reliable expert assistant, providing precise, accurate, and contextually relevant answers to professional physicists' inquiries.
|
|
39
|
+
|
|
40
|
+
## CRITICAL: Source Priority Rules
|
|
41
|
+
**ALWAYS prioritize ATLAS internal sources over web search results for ATLAS-specific technical questions.**
|
|
42
|
+
|
|
43
|
+
## Enhanced Context Information
|
|
44
|
+
You are working with COMBINED SOURCES that include:
|
|
45
|
+
1. **ATLAS Internal Documents** (HIGHEST PRIORITY): Authoritative internal documentation (TWiki, CDS, Indico, AtlasTalk, MkDocs)
|
|
46
|
+
2. **Web Search Results** (SUPPLEMENTARY): Recent external information from the internet (marked as "WebSearch" type)
|
|
47
|
+
|
|
48
|
+
## Core Principles
|
|
49
|
+
1. **Source Hierarchy**: ATLAS internal sources are AUTHORITATIVE for ATLAS-specific information
|
|
50
|
+
2. **Technical Questions**: For software, procedures, and ATLAS-specific topics, prioritize internal sources
|
|
51
|
+
3. **Web Sources**: Use only for recent developments, external context, or when internal sources lack information
|
|
52
|
+
4. **Source Fidelity**: Base all responses exclusively on provided context
|
|
53
|
+
5. **Scientific Rigor**: Maintain highest standards of technical accuracy
|
|
54
|
+
|
|
55
|
+
## Response Format
|
|
56
|
+
Your response must follow this JSON structure:
|
|
57
|
+
{{
|
|
58
|
+
"text": "Your main answer",
|
|
59
|
+
"sources": [
|
|
60
|
+
{{"name": "source_name_1", "utility": "low|medium|high"}},
|
|
61
|
+
{{"name": "source_name_2", "utility": "low|medium|high"}},
|
|
62
|
+
...
|
|
63
|
+
]
|
|
64
|
+
}}
|
|
65
|
+
When a source name contains quotation marks, omit the quotation marks in your response.
|
|
66
|
+
|
|
67
|
+
## Enhanced Context Management Guidelines
|
|
68
|
+
|
|
69
|
+
### MANDATORY Source Priority Rules:
|
|
70
|
+
- **For ATLAS Software/Tools Questions**: ALWAYS prioritize internal ATLAS documentation
|
|
71
|
+
- **For Technical Procedures**: Use internal sources as primary reference
|
|
72
|
+
- **For Recent Updates**: Cross-reference internal sources with web sources
|
|
73
|
+
- **When Internal Sources Available**: Use web sources only for supplementary context
|
|
74
|
+
|
|
75
|
+
### Source Type Handling:
|
|
76
|
+
- **ATLAS Internal Sources** (source_priority: 1):
|
|
77
|
+
* AUTHORITATIVE for all ATLAS-specific information
|
|
78
|
+
* Use as PRIMARY source for software, procedures, detector specifications
|
|
79
|
+
* Trust over external sources for collaboration decisions and technical details
|
|
80
|
+
|
|
81
|
+
- **Web Search Results** (source_priority: 2):
|
|
82
|
+
* Use ONLY when internal sources are insufficient
|
|
83
|
+
* Good for recent developments and external perspectives
|
|
84
|
+
* ALWAYS indicate when information comes from external web sources
|
|
85
|
+
* NEVER contradict internal sources without explicit justification
|
|
86
|
+
|
|
87
|
+
### Information Integration Strategy:
|
|
88
|
+
1. **Primary Response**: Base answer on ATLAS internal sources
|
|
89
|
+
2. **Supplementary Information**: Add web sources only if they provide additional value
|
|
90
|
+
3. **Conflict Resolution**: ALWAYS prioritize ATLAS internal sources for ATLAS-specific matters
|
|
91
|
+
4. **Source Attribution**: Clearly distinguish between internal and external sources
|
|
92
|
+
|
|
93
|
+
### Source Citation Requirements:
|
|
94
|
+
- Clearly distinguish between internal ATLAS sources and web search results
|
|
95
|
+
- When using web sources, mention their external nature in your response
|
|
96
|
+
- Rate utility based on both relevance and source reliability
|
|
97
|
+
- Flag any discrepancies between internal and external sources
|
|
98
|
+
|
|
99
|
+
## Response Requirements
|
|
100
|
+
|
|
101
|
+
### Answer Construction Priority:
|
|
102
|
+
1. **First**: Extract information from ATLAS internal sources
|
|
103
|
+
2. **Second**: Supplement with relevant web search findings ONLY if needed
|
|
104
|
+
3. **Always**: Indicate source type in your response
|
|
105
|
+
4. **Never**: Let web sources override internal ATLAS documentation
|
|
106
|
+
|
|
107
|
+
### Structure and Content
|
|
108
|
+
1. Technical Precision
|
|
109
|
+
- Use ATLAS-standard terminology
|
|
110
|
+
- Follow HEP unit conventions
|
|
111
|
+
- Maintain consistent notation for particles and processes
|
|
112
|
+
- Indicate when information comes from external vs. internal sources
|
|
113
|
+
|
|
114
|
+
2. Clarity Elements
|
|
115
|
+
- Break complex topics into logical steps
|
|
116
|
+
- Use Markdown formatting for readability
|
|
117
|
+
- Structure using headers, lists, and tables when appropriate
|
|
118
|
+
- Clearly separate internal ATLAS information from external web information
|
|
119
|
+
|
|
120
|
+
3. Enhanced Metadata Integration
|
|
121
|
+
- Include relevant run numbers from internal sources
|
|
122
|
+
- Specify √s values when applicable
|
|
123
|
+
- Note integrated luminosity from official sources
|
|
124
|
+
- Reference data-taking periods
|
|
125
|
+
- Mention software/calibration versions
|
|
126
|
+
- Include recent developments from web sources when relevant
|
|
127
|
+
|
|
128
|
+
### Multi-Source Analysis
|
|
129
|
+
1. Identify information from ATLAS internal sources
|
|
130
|
+
2. Supplement with relevant web search findings ONLY if needed
|
|
131
|
+
3. Synthesize comprehensive understanding
|
|
132
|
+
4. Present balanced view with source attribution
|
|
133
|
+
5. Highlight any temporal considerations or recent updates
|
|
134
|
+
|
|
135
|
+
### Context
|
|
136
|
+
{context}
|
|
137
|
+
|
|
138
|
+
### Question
|
|
139
|
+
{question}
|
|
140
|
+
|
|
141
|
+
### Response Requirements
|
|
142
|
+
Provide your response prioritizing ATLAS internal sources:
|
|
143
|
+
{{
|
|
144
|
+
"text": "Your answer prioritizing internal ATLAS sources, supplemented by web sources only when necessary",
|
|
145
|
+
"sources": [
|
|
146
|
+
{{"name": "source_name_1", "utility": "low|medium|high"}},
|
|
147
|
+
{{"name": "source_name_2", "utility": "low|medium|high"}},
|
|
148
|
+
...
|
|
149
|
+
]
|
|
150
|
+
}}
|
|
151
|
+
"""
|
|
File without changes
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
from langchain_core.runnables import RunnableLambda, RunnableParallel, RunnablePassthrough, RunnableSequence
|
|
2
|
+
|
|
3
|
+
from chATLAS_Embed.Base import VectorStore
|
|
4
|
+
from chATLAS_Embed.LangChainVectorStore import LangChainVectorStore
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def search_runnable(
|
|
8
|
+
vectorstore: VectorStore | list[VectorStore],
|
|
9
|
+
) -> RunnableSequence:
|
|
10
|
+
"""
|
|
11
|
+
LangChain RunnableSequence to search one or more vectorstores in parallel.
|
|
12
|
+
|
|
13
|
+
:param vectorstore: A single `VectorStore` or a list of `VectorStore` instances
|
|
14
|
+
:type vectorstore: VectorStore or list[VectorStore]
|
|
15
|
+
|
|
16
|
+
:return: A LangChain `RunnableSequence` that performs retrieval. Results from multiple vectorstores are merged into a single list.
|
|
17
|
+
:rtype: RunnableSequence
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
# Create a list of retrievers from each vectorstore
|
|
21
|
+
if isinstance(vectorstore, list):
|
|
22
|
+
retrievers = [LangChainVectorStore(vector_store=vs) for vs in vectorstore]
|
|
23
|
+
else:
|
|
24
|
+
retrievers = [LangChainVectorStore(vector_store=vectorstore)]
|
|
25
|
+
|
|
26
|
+
# Create parallel retrieval for each retriever
|
|
27
|
+
retrieved_documents = RunnableParallel(
|
|
28
|
+
{f"docs_{i}": retriever for i, retriever in enumerate(retrievers)} | {"question": RunnablePassthrough()}
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
# Merge all retrieved documents into a single list
|
|
32
|
+
def merge_docs(x):
|
|
33
|
+
all_docs = []
|
|
34
|
+
for i in range(len(retrievers)):
|
|
35
|
+
all_docs.extend(x[f"docs_{i}"])
|
|
36
|
+
return all_docs
|
|
37
|
+
|
|
38
|
+
# take the retrived docs and merge them, also pass through the question
|
|
39
|
+
processed = RunnableParallel(
|
|
40
|
+
{"docs": RunnableLambda(merge_docs), "question": RunnableLambda(lambda x: x["question"])}
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
searcher = retrieved_documents | processed
|
|
44
|
+
|
|
45
|
+
return searcher
|
|
File without changes
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
from langchain_core.documents import Document
|
|
2
|
+
from langchain_core.prompts import PromptTemplate, format_document
|
|
3
|
+
|
|
4
|
+
from chATLAS_Chains.prompt.doc_joiners import DEFAULT_DOCUMENT_JOINER
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def combine_documents(
|
|
8
|
+
docs: list[Document], document_prompt: str = DEFAULT_DOCUMENT_JOINER, document_separator: str = "\n\n"
|
|
9
|
+
):
|
|
10
|
+
"""
|
|
11
|
+
Combine a list of documents into a single formatted string.
|
|
12
|
+
|
|
13
|
+
:param docs: The list of documents to combine.
|
|
14
|
+
:type docs: list[Document]
|
|
15
|
+
:param document_prompt: The prompt template used to format each document.
|
|
16
|
+
Defaults to `DEFAULT_DOCUMENT_JOINER`.
|
|
17
|
+
:type document_prompt: str, optional
|
|
18
|
+
:param document_separator: The separator to place between documents in the final string.
|
|
19
|
+
Defaults to two newlines.
|
|
20
|
+
:type document_separator: str, optional
|
|
21
|
+
|
|
22
|
+
:return: A single string containing all formatted documents joined by the separator.
|
|
23
|
+
:rtype: str
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
doc_strings = [format_document(doc, PromptTemplate.from_template(document_prompt)) for doc in docs]
|
|
27
|
+
|
|
28
|
+
return document_separator.join(doc_strings)
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import socket
|
|
3
|
+
|
|
4
|
+
from chATLAS_Chains.log import logger
|
|
5
|
+
from chATLAS_Embed.EmbeddingModels import SentenceTransformerEmbedding
|
|
6
|
+
from chATLAS_Embed.VectorStores import PostgresParentChildVectorStore
|
|
7
|
+
|
|
8
|
+
# Keeps track of which db is stored where
|
|
9
|
+
db_config = {
|
|
10
|
+
"cds_prod_v1": {"hostname": "dbod-chatlas-clone-20250326103412.cern.ch", "port": 6602},
|
|
11
|
+
"twiki_prod": {"hostname": "dbod-chatlas.cern.ch", "port": 6624},
|
|
12
|
+
"indico_prod_v1": {"hostname": "dbod-chatlas.cern.ch", "port": 6624},
|
|
13
|
+
}
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def check_port_forwarding(host="localhost.cern.ch", port=6624):
|
|
17
|
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
18
|
+
result = s.connect_ex((host, port)) # Returns 0 if connection succeeds
|
|
19
|
+
if result == 0:
|
|
20
|
+
logger.debug(f"Port {port} is forwarded and accessible.")
|
|
21
|
+
return True
|
|
22
|
+
else:
|
|
23
|
+
logger.error(f"Port {port} is NOT accessible.")
|
|
24
|
+
return False
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def initialize_vector_store(model_path, postgres_password, db_name):
|
|
28
|
+
"""
|
|
29
|
+
Initialises the connecton to the vector store
|
|
30
|
+
|
|
31
|
+
Args:
|
|
32
|
+
postgres_password: password, usually set at CHATLAS_DB_PASSWORD
|
|
33
|
+
|
|
34
|
+
Returns:
|
|
35
|
+
vector_store: The vector store object
|
|
36
|
+
"""
|
|
37
|
+
logger.info(f"Instantiating vector store {db_name}")
|
|
38
|
+
|
|
39
|
+
# Define the embedding model
|
|
40
|
+
embedding_model = SentenceTransformerEmbedding(model_name=model_path)
|
|
41
|
+
|
|
42
|
+
if os.getenv("CHATLAS_PORT_FORWARDING"):
|
|
43
|
+
logger.debug("Checking port forwarding")
|
|
44
|
+
server_name = "localhost.cern.ch"
|
|
45
|
+
|
|
46
|
+
# check the port is setup properly
|
|
47
|
+
port = db_config[db_name]["port"]
|
|
48
|
+
if not check_port_forwarding(host=server_name, port=port):
|
|
49
|
+
raise ValueError(f"Port forwarding not set up for {port}")
|
|
50
|
+
else:
|
|
51
|
+
logger.debug("Port forwarding not set up, connecting to DB directly")
|
|
52
|
+
server_name = db_config[db_name]["hostname"]
|
|
53
|
+
port = db_config[db_name]["port"]
|
|
54
|
+
|
|
55
|
+
connection_string = f"postgresql://admin:{postgres_password}@{server_name}:{port}/{db_name}"
|
|
56
|
+
|
|
57
|
+
# Connect to the vector store
|
|
58
|
+
vectorstore = PostgresParentChildVectorStore(connection_string=connection_string, embedding_model=embedding_model)
|
|
59
|
+
# Manually set search hyperparams
|
|
60
|
+
vectorstore.search_hyperparams = """
|
|
61
|
+
SET plan_cache_mode = force_generic_plan;
|
|
62
|
+
"""
|
|
63
|
+
|
|
64
|
+
# vectorstore.explain_analyse = True
|
|
65
|
+
logger.info("Vector store initialised")
|
|
66
|
+
|
|
67
|
+
return vectorstore
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
# ADDITIONAL INFO
|
|
71
|
+
# For vectorstores can set with_timings=True to see how long queries are taking on backend for each search
|
|
72
|
+
# Can also set vectorstore.explain_analyse = True to return the EXPLAIN ANALYSE of a query every time before running
|
|
73
|
+
# which can be used to see what in query is taking the most time.
|
|
74
|
+
|
|
75
|
+
# initialize vector store
|
|
76
|
+
if os.getenv("CHATLAS_DB_PASSWORD") is None:
|
|
77
|
+
raise ValueError("CHATLAS_DB_PASSWORD not set in environment")
|
|
78
|
+
|
|
79
|
+
if os.getenv("CHATLAS_EMBEDDING_MODEL_PATH") is None:
|
|
80
|
+
raise ValueError("CHATLAS_EMBEDDING_MODEL_PATH not set in environment")
|
|
81
|
+
else:
|
|
82
|
+
embedding_model_path = os.getenv("CHATLAS_EMBEDDING_MODEL_PATH")
|
|
83
|
+
|
|
84
|
+
twiki_vectorstore = initialize_vector_store(
|
|
85
|
+
model_path=embedding_model_path, postgres_password=os.getenv("CHATLAS_DB_PASSWORD"), db_name="twiki_prod"
|
|
86
|
+
)
|
|
87
|
+
cds_vectorstore = initialize_vector_store(
|
|
88
|
+
model_path=embedding_model_path, postgres_password=os.getenv("CHATLAS_DB_PASSWORD"), db_name="cds_prod_v1"
|
|
89
|
+
)
|
|
90
|
+
indico_vectorstore = initialize_vector_store(
|
|
91
|
+
model_path=embedding_model_path, postgres_password=os.getenv("CHATLAS_DB_PASSWORD"), db_name="indico_prod_v1"
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
vectorstore = [twiki_vectorstore, cds_vectorstore, indico_vectorstore]
|
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: chATLAS_Chains
|
|
3
|
+
Version: 0.1.3
|
|
4
|
+
Summary: A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project.
|
|
5
|
+
Author-email: Joe Egan <joseph.caimin.egan@cern.ch>
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://gitlab.cern.ch/belliot/chatlas-packages/
|
|
8
|
+
Project-URL: Documentation, https://chatlas-packages.docs.cern.ch/chATLAS_Chain/
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Requires-Python: >=3.11
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
License-File: LICENSE
|
|
15
|
+
Requires-Dist: chATLAS_Benchmark>=0.0.9
|
|
16
|
+
Requires-Dist: chATLAS_Embed>=0.1.14
|
|
17
|
+
Requires-Dist: langchain~=0.3.3
|
|
18
|
+
Requires-Dist: langchain_core
|
|
19
|
+
Requires-Dist: langchain_openai
|
|
20
|
+
Dynamic: license-file
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
# chATLAS_Chains
|
|
24
|
+
|
|
25
|
+
This package implements and benchmarks various Retrieval Augmented Generation (RAG) chains for use in the [chATLAS](https://chatlas-flask-chatlas.app.cern.ch) project.
|
|
26
|
+
|
|
27
|
+
## Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
conda create -n venv chatlas_chains_env python=3.10
|
|
31
|
+
conda activate chatlas_chains_env
|
|
32
|
+
pip install chatlas-chains
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## Environment variables
|
|
36
|
+
|
|
37
|
+
These are required for the following use cases
|
|
38
|
+
|
|
39
|
+
1. Using an OpenAI LLM
|
|
40
|
+
```bash
|
|
41
|
+
export CHATLAS_OPENAI_KEY='your api key'
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
2. Benchmarking, set the path to the question set
|
|
45
|
+
```bash
|
|
46
|
+
export CHATLAS_BENCHMARK_QUESTIONS=/path/to/questions.josn
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
## Available Chains
|
|
50
|
+
- chains.basic.basic_retrieval_chain
|
|
51
|
+
- chains.basic_graph.basic_retrieval_graph
|
|
52
|
+
|
|
53
|
+
## Benchmarking
|
|
54
|
+
|
|
55
|
+
To benchmark e.g. the chains in `chATLAS_Chains.chains.basic` run this from the project root
|
|
56
|
+
```bash
|
|
57
|
+
python benchmark/basic.py
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
## Testing
|
|
61
|
+
|
|
62
|
+
The tests require a running postgres server to work. If on lxplus you can modify `TEST_DB_CONFIG` in [tests/conftest.py](tests/conftest.py) to connect to the chATLAS server.
|
|
63
|
+
|
|
64
|
+
If you want to create a local dummy postgres server, you need to install `psql`. This can be done on macOS using [homebrew](https://brew.sh):
|
|
65
|
+
|
|
66
|
+
Software install
|
|
67
|
+
```bash
|
|
68
|
+
brew install postgresql
|
|
69
|
+
brew services start postgresql
|
|
70
|
+
brew install pgvector
|
|
71
|
+
brew unlink pgvector && brew link pgvector
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
Create a user
|
|
75
|
+
```bash
|
|
76
|
+
psql -h localhost -U postgres
|
|
77
|
+
ALTER USER postgres WITH PASSWORD 'Set_your_password_here';
|
|
78
|
+
CREATE EXTENSION IF NOT EXISTS vector;
|
|
79
|
+
```
|
|
80
|
+
## CHANGELOG
|
|
81
|
+
|
|
82
|
+
#### 0.1.3
|
|
83
|
+
|
|
84
|
+
Fixing imports
|
|
85
|
+
|
|
86
|
+
Changed output format of `basic_retrieval_chain` (`docs` key is now a list of `Document` objects, rather than a dict)
|
|
87
|
+
|
|
88
|
+
Unit tests for `basic_retrieval_chain`
|
|
89
|
+
|
|
90
|
+
#### 0.1.2
|
|
91
|
+
|
|
92
|
+
Unit tests
|
|
93
|
+
|
|
94
|
+
First Langgraph chain
|
|
95
|
+
|
|
96
|
+
#### 0.1.1
|
|
97
|
+
|
|
98
|
+
Initial Release
|
|
99
|
+
|
|
100
|
+
---
|
|
101
|
+
## 📄 License
|
|
102
|
+
|
|
103
|
+
chATLAS_Benchmark is released under Apache v2.0 license.
|
|
104
|
+
|
|
105
|
+
---
|
|
106
|
+
|
|
107
|
+
<div align="center">
|
|
108
|
+
|
|
109
|
+
**Made with ❤️ by the ATLAS Collaboration**
|
|
110
|
+
|
|
111
|
+
*For questions and support, please [contact](mailto:joseph.caimin.egan@cern.ch)*
|
|
112
|
+
|
|
113
|
+
</div>
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
benchmark/basic.py,sha256=9A0MGfk_0UTKiIuqexRSiyGwt1OY1fErhiTmY_7AVCE,1257
|
|
2
|
+
benchmark/conversational.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
3
|
+
chATLAS_Chains/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
4
|
+
chATLAS_Chains/log.py,sha256=8LuEBDKnVr-I6bhr-LDnpemcPNuneGGq7_DS9vOCovU,649
|
|
5
|
+
chATLAS_Chains/vectorstore.py,sha256=2VOZGTPVmeRNaxMY2GpaQYQKzMQ9TtwAWcw9JwBna9Q,3630
|
|
6
|
+
chATLAS_Chains/chains/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
7
|
+
chATLAS_Chains/chains/basic.py,sha256=AiBhbHTNuofrQECyEzKqXX2bBhdN4zi77kce2VQgWFA,1403
|
|
8
|
+
chATLAS_Chains/chains/basic_graph.py,sha256=bqFtC8BLl8zJNGe2EuHPnRXK9Pm2ya3YstY7quPiFR4,5324
|
|
9
|
+
chATLAS_Chains/chains/conversational.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
10
|
+
chATLAS_Chains/chains/websearch_retrieval_chain.py,sha256=lAvy2ags4dDSN2inOtVhctQoJ10z2bIv2chkTxATGFs,15330
|
|
11
|
+
chATLAS_Chains/llm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
12
|
+
chATLAS_Chains/llm/model_selection.py,sha256=FzJ1I4x2E4QzPWjSQIye0kcEEDdovB7_DeIIQK43QjY,1648
|
|
13
|
+
chATLAS_Chains/prompt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
14
|
+
chATLAS_Chains/prompt/doc_joiners.py,sha256=yvvK3B8VUszGQZJzx4WUaYyMzVJdV96FRsv8AoA7Bkc,133
|
|
15
|
+
chATLAS_Chains/prompt/starters.py,sha256=98HhjagX5OjwARN5QP4zhywjUp4gTuHmIKTknvSTMZ8,9609
|
|
16
|
+
chATLAS_Chains/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
17
|
+
chATLAS_Chains/search/basic.py,sha256=FmwGSAKZGE7Jd0DhVId3qQIiHDqInGXh6nuGfecHJSo,1685
|
|
18
|
+
chATLAS_Chains/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
19
|
+
chATLAS_Chains/utils/doc_utils.py,sha256=wgyw4t7O41egF0OyoHugDyMO4vjjbUqu_aFehebD0G0,1122
|
|
20
|
+
chatlas_chains-0.1.3.dist-info/licenses/LICENSE,sha256=i4xkIjVh9qAZlkujBgWWzSWMeV_dv9trQtyiiT18p5c,11381
|
|
21
|
+
tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
22
|
+
tests/conftest.py,sha256=UxMU8nsteKKfMvdsqbkU4T9b58_kwGr3_GsWa19F79E,13324
|
|
23
|
+
tests/test_chains.py,sha256=ZLAbIPeF231ONTi3mIkKPcJa6wPH9P3tWMida1Sempw,1815
|
|
24
|
+
tests/test_llm.py,sha256=oeyTxX4eZLgIOSRd3UdML_VPq5SK4U23rpfmUZar6N8,3168
|
|
25
|
+
tests/test_search.py,sha256=3Mrjuip5OprjzfgkUZ_XV8ljSeBiN-tkAGYyAvwqSMo,1465
|
|
26
|
+
tests/test_utils.py,sha256=5r5mGm3raEghNTESCcwWmFkYfm5H66AHw5MigCBbjQ4,853
|
|
27
|
+
chatlas_chains-0.1.3.dist-info/METADATA,sha256=-ok82d1hdwlwmTiDCkyUfQL3rQFtg7FmOST9eaVwSLE,2845
|
|
28
|
+
chatlas_chains-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
|
29
|
+
chatlas_chains-0.1.3.dist-info/top_level.txt,sha256=hP2-U0jKBWvH-UebH-3XlDTSJ6BjRWS6HJi8PW7rJkU,31
|
|
30
|
+
chatlas_chains-0.1.3.dist-info/RECORD,,
|