참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
@@ -0,0 +1,4 @@
|
||||
from .graphrag import GraphRAG
|
||||
from .prompts import PromptTemplate, RagTemplate, SchemaExtractionTemplate
|
||||
|
||||
__all__ = ["GraphRAG", "PromptTemplate", "RagTemplate", "SchemaExtractionTemplate"]
|
||||
@@ -0,0 +1,256 @@
|
||||
# Copyright (c) "Neo4j"
|
||||
# Neo4j Sweden AB [https://neo4j.com]
|
||||
# #
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
# #
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
# #
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# built-in dependencies
|
||||
from __future__ import annotations
|
||||
import logging
|
||||
import warnings
|
||||
from typing import Any, List, Optional, Union
|
||||
|
||||
# 3rd party dependencies
|
||||
from pydantic import ValidationError
|
||||
|
||||
# project dependencies
|
||||
from neo4j_graphrag.exceptions import (
|
||||
RagInitializationError,
|
||||
SearchValidationError,
|
||||
)
|
||||
from neo4j_graphrag.generation.prompts import RagTemplate
|
||||
from neo4j_graphrag.generation.types import RagInitModel, RagResultModel, RagSearchModel
|
||||
from neo4j_graphrag.llm import LLMInterface, LLMInterfaceV2
|
||||
from neo4j_graphrag.llm.utils import legacy_inputs_to_messages
|
||||
from neo4j_graphrag.message_history import MessageHistory
|
||||
from neo4j_graphrag.retrievers.base import Retriever
|
||||
from neo4j_graphrag.types import LLMMessage, RetrieverResult
|
||||
from neo4j_graphrag.utils.logging import prettify
|
||||
|
||||
# Set up logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# pylint: disable=raise-missing-from
|
||||
class GraphRAG:
|
||||
"""Performs a GraphRAG search using a specific retriever
|
||||
and LLM.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import neo4j
|
||||
from neo4j_graphrag.retrievers import VectorRetriever
|
||||
from neo4j_graphrag.llm.openai_llm import OpenAILLM
|
||||
from neo4j_graphrag.generation import GraphRAG
|
||||
|
||||
driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
|
||||
|
||||
retriever = VectorRetriever(driver, "vector-index-name", custom_embedder)
|
||||
llm = OpenAILLM()
|
||||
graph_rag = GraphRAG(retriever, llm)
|
||||
graph_rag.search(query_text="Find me a book about Fremen")
|
||||
|
||||
Args:
|
||||
retriever (Retriever): The retriever used to find relevant context to pass to the LLM.
|
||||
llm (LLMInterface, LLMInterfaceV2 or LangChain Chat Model): The LLM used to generate
|
||||
the answer.
|
||||
prompt_template (RagTemplate): The prompt template that will be formatted with context and
|
||||
user question and passed to the LLM.
|
||||
|
||||
Raises:
|
||||
RagInitializationError: If validation of the input arguments fail.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
retriever: Retriever,
|
||||
llm: Union[LLMInterface, LLMInterfaceV2, Any],
|
||||
prompt_template: RagTemplate = RagTemplate(),
|
||||
):
|
||||
try:
|
||||
validated_data = RagInitModel(
|
||||
retriever=retriever,
|
||||
llm=llm,
|
||||
prompt_template=prompt_template,
|
||||
)
|
||||
except ValidationError as e:
|
||||
raise RagInitializationError(e.errors())
|
||||
self.retriever = validated_data.retriever
|
||||
self.llm = validated_data.llm
|
||||
self.prompt_template = validated_data.prompt_template
|
||||
|
||||
def search(
|
||||
self,
|
||||
query_text: str = "",
|
||||
message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
|
||||
examples: str = "",
|
||||
retriever_config: Optional[dict[str, Any]] = None,
|
||||
return_context: Optional[bool] = None,
|
||||
response_fallback: Optional[str] = None,
|
||||
) -> RagResultModel:
|
||||
"""
|
||||
.. warning::
|
||||
The default value of 'return_context' will change from 'False'
|
||||
to 'True' in a future version.
|
||||
|
||||
|
||||
This method performs a full RAG search:
|
||||
1. Retrieval: context retrieval
|
||||
2. Augmentation: prompt formatting
|
||||
3. Generation: answer generation with LLM
|
||||
|
||||
|
||||
Args:
|
||||
query_text (str): The user question.
|
||||
message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection
|
||||
of previous messages, with each message having a specific role assigned.
|
||||
examples (str): Examples added to the LLM prompt.
|
||||
retriever_config (Optional[dict]): Parameters passed to the retriever.
|
||||
search method; e.g.: top_k
|
||||
return_context (bool): Whether to append the retriever result to the final result
|
||||
(default: False).
|
||||
response_fallback (Optional[str]): If not null, will return this message instead
|
||||
of calling the LLM if context comes back empty.
|
||||
|
||||
Returns:
|
||||
RagResultModel: The LLM-generated answer.
|
||||
|
||||
"""
|
||||
if return_context is None:
|
||||
warnings.warn(
|
||||
"The default value of 'return_context' will change from 'False'"
|
||||
" to 'True' in a future version.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
return_context = False
|
||||
|
||||
try:
|
||||
validated_data = RagSearchModel(
|
||||
query_text=query_text,
|
||||
examples=examples,
|
||||
retriever_config=retriever_config or {},
|
||||
return_context=return_context,
|
||||
response_fallback=response_fallback,
|
||||
)
|
||||
except ValidationError as e:
|
||||
raise SearchValidationError(e.errors())
|
||||
if isinstance(message_history, MessageHistory):
|
||||
message_history = message_history.messages
|
||||
query = self._build_query(validated_data.query_text, message_history)
|
||||
retriever_result: RetrieverResult = self.retriever.search(
|
||||
query_text=query, **validated_data.retriever_config
|
||||
)
|
||||
if len(retriever_result.items) == 0 and response_fallback is not None:
|
||||
answer = response_fallback
|
||||
else:
|
||||
context = "\n".join(item.content for item in retriever_result.items)
|
||||
prompt = self.prompt_template.format(
|
||||
query_text=query_text, context=context, examples=validated_data.examples
|
||||
)
|
||||
|
||||
logger.debug("RAG: retriever_result=%s", prettify(retriever_result))
|
||||
logger.debug("RAG: prompt=%s", prompt)
|
||||
|
||||
if self.is_langchain_compatible():
|
||||
# llm interface v2 or langchain chat model
|
||||
messages = legacy_inputs_to_messages(
|
||||
prompt=prompt,
|
||||
message_history=message_history,
|
||||
system_instruction=self.prompt_template.system_instructions,
|
||||
)
|
||||
|
||||
# langchain chat model compatible invoke
|
||||
llm_response = self.llm.invoke(
|
||||
input=messages,
|
||||
)
|
||||
elif isinstance(self.llm, LLMInterface):
|
||||
# may have custom LLMs inherited from V1, keep it for backward compatibility
|
||||
llm_response = self.llm.invoke(
|
||||
input=prompt,
|
||||
message_history=message_history,
|
||||
system_instruction=self.prompt_template.system_instructions,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Type {type(self.llm)} of LLM is not supported.")
|
||||
answer = llm_response.content
|
||||
result: dict[str, Any] = {"answer": answer}
|
||||
if return_context:
|
||||
result["retriever_result"] = retriever_result
|
||||
return RagResultModel(**result)
|
||||
|
||||
def _build_query(
|
||||
self,
|
||||
query_text: str,
|
||||
message_history: Optional[List[LLMMessage]] = None,
|
||||
) -> str:
|
||||
"""Builds the final query text, incorporating message history if provided."""
|
||||
summary_system_message = (
|
||||
"You are a summarization assistant. "
|
||||
"Summarize the given text in no more than 300 words."
|
||||
)
|
||||
if message_history:
|
||||
summarization_prompt = self._chat_summary_prompt(
|
||||
message_history=message_history
|
||||
)
|
||||
if self.is_langchain_compatible():
|
||||
messages = legacy_inputs_to_messages(
|
||||
summarization_prompt,
|
||||
system_instruction=summary_system_message,
|
||||
)
|
||||
summary = self.llm.invoke(
|
||||
input=messages,
|
||||
).content
|
||||
elif isinstance(self.llm, LLMInterface):
|
||||
summary = self.llm.invoke(
|
||||
input=summarization_prompt,
|
||||
system_instruction=summary_system_message,
|
||||
).content
|
||||
else:
|
||||
raise ValueError(f"Type {type(self.llm)} of LLM is not supported.")
|
||||
|
||||
return self.conversation_prompt(summary=summary, current_query=query_text)
|
||||
return query_text
|
||||
|
||||
def is_langchain_compatible(self) -> bool:
|
||||
"""Checks if the LLM is compatible with LangChain."""
|
||||
if isinstance(self.llm, LLMInterfaceV2):
|
||||
return True
|
||||
|
||||
try:
|
||||
# langchain-core is an optional dependency
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
|
||||
return isinstance(self.llm, BaseChatModel)
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def _chat_summary_prompt(self, message_history: List[LLMMessage]) -> str:
|
||||
message_list = [
|
||||
f"{message['role']}: {message['content']}" for message in message_history
|
||||
]
|
||||
history = "\n".join(message_list)
|
||||
return f"""
|
||||
Summarize the message history:
|
||||
|
||||
{history}
|
||||
"""
|
||||
|
||||
def conversation_prompt(self, summary: str, current_query: str) -> str:
|
||||
return f"""
|
||||
Message Summary:
|
||||
{summary}
|
||||
|
||||
Current Query:
|
||||
{current_query}
|
||||
"""
|
||||
@@ -0,0 +1,314 @@
|
||||
# Copyright (c) "Neo4j"
|
||||
# Neo4j Sweden AB [https://neo4j.com]
|
||||
# #
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
# #
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
# #
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import Any, Optional
|
||||
|
||||
from neo4j_graphrag.exceptions import (
|
||||
PromptMissingInputError,
|
||||
PromptMissingPlaceholderError,
|
||||
)
|
||||
|
||||
|
||||
class PromptTemplate:
|
||||
"""This class is used to generate a parameterized prompt. It is defined
|
||||
from a string (the template) using the Python format syntax (parameters
|
||||
between curly braces `{}`) and a list of required inputs.
|
||||
Before sending the instructions to an LLM, call the `format` method that will
|
||||
replace parameters with the provided values. If any of the expected inputs is
|
||||
missing, a `PromptMissingInputError` is raised.
|
||||
"""
|
||||
|
||||
DEFAULT_SYSTEM_INSTRUCTIONS: str = ""
|
||||
DEFAULT_TEMPLATE: str = ""
|
||||
EXPECTED_INPUTS: list[str] = list()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
template: Optional[str] = None,
|
||||
expected_inputs: Optional[list[str]] = None,
|
||||
system_instructions: Optional[str] = None,
|
||||
) -> None:
|
||||
self.template = template or self.DEFAULT_TEMPLATE
|
||||
self.expected_inputs = expected_inputs or self.EXPECTED_INPUTS
|
||||
self.system_instructions = (
|
||||
system_instructions or self.DEFAULT_SYSTEM_INSTRUCTIONS
|
||||
)
|
||||
|
||||
for e in self.expected_inputs:
|
||||
if f"{{{e}}}" not in self.template:
|
||||
raise PromptMissingPlaceholderError(
|
||||
f"`template` is missing placeholder {e}"
|
||||
)
|
||||
|
||||
def _format(self, **kwargs: Any) -> str:
|
||||
for e in self.EXPECTED_INPUTS:
|
||||
if e not in kwargs:
|
||||
raise PromptMissingInputError(f"Missing input '{e}'")
|
||||
return self.template.format(**kwargs)
|
||||
|
||||
def format(self, *args: Any, **kwargs: Any) -> str:
|
||||
"""This method is used to replace parameters with the provided values.
|
||||
Parameters must be provided:
|
||||
- as kwargs
|
||||
- as args if using the same order as in the expected inputs
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
prompt_template = PromptTemplate(
|
||||
template='''Explain the following concept to {target_audience}:
|
||||
Concept: {concept}
|
||||
Answer:
|
||||
''',
|
||||
expected_inputs=['target_audience', 'concept']
|
||||
)
|
||||
prompt = prompt_template.format('12 yo children', concept='graph database')
|
||||
print(prompt)
|
||||
|
||||
# Result:
|
||||
# '''Explain the following concept to 12 yo children:
|
||||
# Concept: graph database
|
||||
# Answer:
|
||||
# '''
|
||||
|
||||
"""
|
||||
data = kwargs
|
||||
data.update({k: v for k, v in zip(self.expected_inputs, args)})
|
||||
return self._format(**data)
|
||||
|
||||
|
||||
class RagTemplate(PromptTemplate):
|
||||
DEFAULT_SYSTEM_INSTRUCTIONS = "Answer the user question using the provided context."
|
||||
DEFAULT_TEMPLATE = """Context:
|
||||
{context}
|
||||
|
||||
Examples:
|
||||
{examples}
|
||||
|
||||
Question:
|
||||
{query_text}
|
||||
|
||||
Answer:
|
||||
"""
|
||||
EXPECTED_INPUTS = ["context", "query_text", "examples"]
|
||||
|
||||
def format(self, query_text: str, context: str, examples: str) -> str:
|
||||
return super().format(query_text=query_text, context=context, examples=examples)
|
||||
|
||||
|
||||
class Text2CypherTemplate(PromptTemplate):
|
||||
DEFAULT_TEMPLATE = """
|
||||
Task: Generate a Cypher statement for querying a Neo4j graph database from a user input.
|
||||
|
||||
Schema:
|
||||
{schema}
|
||||
|
||||
Examples (optional):
|
||||
{examples}
|
||||
|
||||
Input:
|
||||
{query_text}
|
||||
|
||||
Do not use any properties or relationships not included in the schema.
|
||||
Do not include triple backticks ``` or any additional text except the generated Cypher statement in your response.
|
||||
|
||||
Cypher query:
|
||||
"""
|
||||
EXPECTED_INPUTS = ["query_text"]
|
||||
|
||||
def format(
|
||||
self,
|
||||
schema: Optional[str] = None,
|
||||
examples: Optional[str] = None,
|
||||
query_text: str = "",
|
||||
query: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
if query is not None:
|
||||
if query_text:
|
||||
warnings.warn(
|
||||
"Both 'query' and 'query_text' are provided, 'query_text' will be used.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
elif isinstance(query, str):
|
||||
warnings.warn(
|
||||
"'query' is deprecated and will be removed in a future version, please use 'query_text' instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
query_text = query
|
||||
|
||||
return super().format(
|
||||
query_text=query_text, schema=schema, examples=examples, **kwargs
|
||||
)
|
||||
|
||||
|
||||
class ERExtractionTemplate(PromptTemplate):
|
||||
DEFAULT_TEMPLATE = """
|
||||
You are a top-tier algorithm designed for extracting
|
||||
information in structured formats to build a knowledge graph.
|
||||
|
||||
Extract the entities (nodes) and specify their type from the following text.
|
||||
Also extract the relationships between these nodes.
|
||||
|
||||
Return result as JSON using the following format:
|
||||
{{"nodes": [ {{"id": "0", "label": "Person", "properties": {{"name": "John"}} }}],
|
||||
"relationships": [{{"type": "KNOWS", "start_node_id": "0", "end_node_id": "1", "properties": {{"since": "2024-08-01"}} }}] }}
|
||||
|
||||
Use only the following node and relationship types (if provided):
|
||||
{schema}
|
||||
|
||||
Assign a unique ID (string) to each node, and reuse it to define relationships.
|
||||
Do respect the source and target node types for relationship and
|
||||
the relationship direction.
|
||||
|
||||
Make sure you adhere to the following rules to produce valid JSON objects:
|
||||
- Do not return any additional information other than the JSON in it.
|
||||
- Omit any backticks around the JSON - simply output the JSON on its own.
|
||||
- The JSON object must not wrapped into a list - it is its own JSON object.
|
||||
- Property names must be enclosed in double quotes
|
||||
|
||||
Examples:
|
||||
{examples}
|
||||
|
||||
Input text:
|
||||
|
||||
{text}
|
||||
"""
|
||||
EXPECTED_INPUTS = ["text"]
|
||||
|
||||
def format(
|
||||
self,
|
||||
schema: dict[str, Any],
|
||||
examples: str,
|
||||
text: str = "",
|
||||
) -> str:
|
||||
return super().format(text=text, schema=schema, examples=examples)
|
||||
|
||||
|
||||
class SchemaExtractionTemplate(PromptTemplate):
|
||||
DEFAULT_TEMPLATE = """
|
||||
You are a top-tier algorithm designed for extracting a labeled property graph schema in
|
||||
structured formats.
|
||||
|
||||
Generate a generalized graph schema based on the input text. Identify key node types,
|
||||
their relationship types, and property types.
|
||||
|
||||
IMPORTANT RULES:
|
||||
1. Return only abstract schema information, not concrete instances.
|
||||
2. Use singular PascalCase labels for node types (e.g., Person, Company, Product).
|
||||
3. Use UPPER_SNAKE_CASE labels for relationship types (e.g., WORKS_FOR, MANAGES).
|
||||
4. Include property definitions only when the type can be confidently inferred, otherwise omit them.
|
||||
5. When defining patterns, ensure that every node label and relationship label mentioned exists in your lists of node types and relationship types.
|
||||
6. Do not create node types that aren't clearly mentioned in the text.
|
||||
7. Keep your schema minimal and focused on clearly identifiable patterns in the text.
|
||||
8. UNIQUENESS CONSTRAINTS (optional, node properties only):
|
||||
8.1 Each node type may have at most one UNIQUENESS constraint in typical designs.
|
||||
8.2 Only use properties that seem to not have too many missing values in the sample.
|
||||
8.3 Constraints reference node_types by label and specify which properties are unique via "property_names" (a list of one or more property names).
|
||||
8.4 Every property in a uniqueness constraint MUST also appear in the corresponding node_type as a property.
|
||||
8.5 Uniqueness does NOT imply that the property must exist on every node (existence is separate; see rule 9).
|
||||
8.6 For composite uniqueness (multiple properties), list all properties in "property_names". The combination of values must be unique.
|
||||
9. EXISTENCE CONSTRAINTS (optional, single property only):
|
||||
9.1 Use EXISTENCE constraints to mark properties that MUST be present (non-null) on every instance.
|
||||
9.2 For a node property, add {{"type": "EXISTENCE", "node_type": "<Label>", "property_names": ["<name>"], "relationship_type": ""}} (use an empty string, not null, when the constraint is not on a relationship).
|
||||
9.3 For a relationship property, add {{"type": "EXISTENCE", "node_type": "", "property_names": ["<name>"], "relationship_type": "<REL_TYPE>"}}.
|
||||
9.4 Each EXISTENCE constraint must reference exactly one of node_type or relationship_type (non-empty), never both.
|
||||
9.5 Do not infer EXISTENCE from UNIQUENESS; they are independent (as in Neo4j Cypher constraints).
|
||||
9.6 EXISTENCE constraints must have exactly one property in "property_names" (composite existence is not supported).
|
||||
10. KEY CONSTRAINTS (optional, Neo4j NODE KEY / RELATIONSHIP KEY):
|
||||
10.1 Use KEY when properties must exist on every instance and together form the natural identifier (uniqueness + mandatory presence).
|
||||
10.2 Same wire shape as EXISTENCE: exactly one of node_type or relationship_type (non-empty), with the other as "".
|
||||
10.3 Do not combine UNIQUENESS and KEY on the same node type with the same properties.
|
||||
10.4 Do not infer KEY from UNIQUENESS alone; KEY implies required presence, unlike UNIQUENESS alone.
|
||||
10.5 For composite key (multiple properties), list all properties in "property_names". All must exist and the combination must be unique.
|
||||
11. Never use double underscores (__) as a prefix or suffix in node labels or relationship types (e.g. __Person__ or __KNOWS__ are forbidden).
|
||||
|
||||
Accepted property types are: BOOLEAN, DATE, DURATION, FLOAT, INTEGER, LIST,
|
||||
LOCAL_DATETIME, LOCAL_TIME, POINT, STRING, ZONED_DATETIME, ZONED_TIME.
|
||||
|
||||
Return a valid JSON object that follows this precise structure:
|
||||
{{
|
||||
"node_types": [
|
||||
{{
|
||||
"label": "Person",
|
||||
"properties": [
|
||||
{{
|
||||
"name": "name",
|
||||
"type": "STRING"
|
||||
}},
|
||||
{{
|
||||
"name": "email",
|
||||
"type": "STRING"
|
||||
}},
|
||||
{{
|
||||
"name": "employee_id",
|
||||
"type": "STRING"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
...
|
||||
],
|
||||
"relationship_types": [
|
||||
{{
|
||||
"label": "WORKS_FOR"
|
||||
}}
|
||||
...
|
||||
],
|
||||
"patterns": [
|
||||
{{"source": "Person", "relationship": "WORKS_FOR", "target": "Company"}},
|
||||
...
|
||||
],
|
||||
"constraints": [
|
||||
{{
|
||||
"type": "UNIQUENESS",
|
||||
"node_type": "Person",
|
||||
"property_names": ["email"],
|
||||
"relationship_type": ""
|
||||
}},
|
||||
{{
|
||||
"type": "EXISTENCE",
|
||||
"node_type": "Person",
|
||||
"property_names": ["name"],
|
||||
"relationship_type": ""
|
||||
}},
|
||||
{{
|
||||
"type": "KEY",
|
||||
"node_type": "Person",
|
||||
"property_names": ["employee_id"],
|
||||
"relationship_type": ""
|
||||
}}
|
||||
...
|
||||
]
|
||||
}}
|
||||
|
||||
Examples:
|
||||
{examples}
|
||||
|
||||
Input text:
|
||||
{text}
|
||||
"""
|
||||
EXPECTED_INPUTS = ["text"]
|
||||
|
||||
def format(
|
||||
self,
|
||||
text: str = "",
|
||||
examples: str = "",
|
||||
) -> str:
|
||||
return super().format(text=text, examples=examples)
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) "Neo4j"
|
||||
# Neo4j Sweden AB [https://neo4j.com]
|
||||
# #
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
# #
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
# #
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, field_validator
|
||||
|
||||
from neo4j_graphrag.generation.prompts import RagTemplate
|
||||
from neo4j_graphrag.retrievers.base import Retriever
|
||||
from neo4j_graphrag.types import RetrieverResult
|
||||
|
||||
|
||||
class RagInitModel(BaseModel):
|
||||
retriever: Retriever
|
||||
llm: Any
|
||||
prompt_template: RagTemplate
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@field_validator("llm")
|
||||
def check_llm(cls, value: Any) -> Any:
|
||||
invoke = getattr(value, "invoke", None)
|
||||
if invoke and callable(invoke):
|
||||
return value
|
||||
raise ValueError("llm must be callable")
|
||||
|
||||
|
||||
class RagSearchModel(BaseModel):
|
||||
query_text: str
|
||||
examples: str = ""
|
||||
retriever_config: dict[str, Any] = {}
|
||||
return_context: bool = False
|
||||
response_fallback: Optional[str] = None
|
||||
|
||||
|
||||
class RagResultModel(BaseModel):
|
||||
answer: str
|
||||
retriever_result: Optional[RetrieverResult] = None
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
Reference in New Issue
Block a user