126 lines
5.3 KiB
Python
126 lines
5.3 KiB
Python
import logging
|
|
from langchain_experimental.graph_transformers import LLMGraphTransformer
|
|
from langchain_openai import ChatOpenAI
|
|
from django.conf import settings
|
|
from langchain_community.graphs import Neo4jGraph
|
|
from langchain.schema import Document
|
|
from .models import Neo4jProfile
|
|
from datetime import datetime
|
|
from langchain.prompts import ChatPromptTemplate
|
|
from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_PROFILE_NAME = "DefaultNeo4jProfile"
|
|
|
|
class Neo4jDatabase:
|
|
"""
|
|
Handles connection to Neo4j using different profiles for various use cases.
|
|
"""
|
|
|
|
def __init__(self, profile_name=None):
|
|
try:
|
|
# Load the specified Neo4j profile, or fallback to the default
|
|
if profile_name:
|
|
self.profile = Neo4jProfile.objects.get(name=profile_name)
|
|
else:
|
|
logger.warning("No profile specified. Using default Neo4j profile.")
|
|
self.profile = Neo4jProfile.objects.get(name=DEFAULT_PROFILE_NAME)
|
|
|
|
# Attempt to connect to the assigned Neo4j instance
|
|
try:
|
|
self.graph = Neo4jGraph(
|
|
url=self.profile.uri,
|
|
username=self.profile.username,
|
|
password=self.profile.password
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize Neo4jGraph: {e}")
|
|
self.graph = None
|
|
|
|
# Initialize LLM if graph is available
|
|
try:
|
|
self.llm = ChatOpenAI(
|
|
api_key=self.profile.openai_api_key,
|
|
model_name=self.profile.model_name,
|
|
temperature=0
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize ChatOpenAI: {e}")
|
|
self.llm = None
|
|
|
|
# Prepare graph transformer only if llm available
|
|
if self.llm:
|
|
english_prompt = ChatPromptTemplate.from_template(
|
|
"Transform the following text into a graph structure. All nodes, relationships, and properties should be in English, regardless of the original language."
|
|
)
|
|
try:
|
|
self.graph_transformer = LLMGraphTransformer(
|
|
llm=self.llm,
|
|
prompt=english_prompt
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize LLMGraphTransformer: {e}")
|
|
self.graph_transformer = None
|
|
else:
|
|
self.graph_transformer = None
|
|
|
|
logger.info(f"Neo4jDatabase initialized with profile: {self.profile.name}")
|
|
|
|
except Neo4jProfile.DoesNotExist:
|
|
logger.error(f"Neo4j profile '{profile_name}' not found.")
|
|
self.graph = None
|
|
self.llm = None
|
|
self.graph_transformer = None
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error initializing Neo4jDatabase: {e}")
|
|
self.graph = None
|
|
self.llm = None
|
|
self.graph_transformer = None
|
|
|
|
def store_interaction(self, user_id, bot_id, user_message, bot_response, platform):
|
|
"""
|
|
Stores a chatbot interaction as a structured graph in Neo4j.
|
|
Converts messages into `Document` objects to work with `LLMGraphTransformer`.
|
|
"""
|
|
if not self.graph_transformer:
|
|
logger.warning("Graph transformer not available, skipping store_interaction.")
|
|
return
|
|
|
|
try:
|
|
timestamp = datetime.utcnow().isoformat()
|
|
documents = [
|
|
Document(page_content=user_message, metadata={"role": "user", "user_id": user_id, "platform": platform, "created_at": timestamp}),
|
|
Document(page_content=bot_response, metadata={"role": "bot", "bot_id": bot_id, "platform": platform, "created_at": timestamp}),
|
|
]
|
|
graph_docs = self.graph_transformer.convert_to_graph_documents(documents)
|
|
if self.graph:
|
|
self.graph.add_graph_documents(graph_docs, include_source=True)
|
|
logger.info(f"Stored interaction in Neo4j (Profile: {self.profile.name})")
|
|
else:
|
|
logger.warning("No Neo4j graph available, skipping add_graph_documents.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to store interaction in Neo4j: {e}")
|
|
|
|
def query_graph(self, user_query):
|
|
"""
|
|
Queries the graph using GraphCypherQAChain and returns a structured response.
|
|
"""
|
|
if not self.graph or not self.llm:
|
|
logger.warning("Graph or LLM not available, cannot query.")
|
|
return None
|
|
|
|
try:
|
|
qa_chain = GraphCypherQAChain.from_llm(
|
|
llm=self.llm,
|
|
graph=self.graph,
|
|
verbose=True,
|
|
allow_dangerous_requests=True
|
|
)
|
|
result = qa_chain.invoke({"query": user_query})
|
|
logger.info(f"Resulting Neo4j query result: {result}")
|
|
return result.get('result')
|
|
except Exception as e:
|
|
logger.error(f"Graph query failed: {e}")
|
|
return None
|