112 lines
4.3 KiB
Python
112 lines
4.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" # Define a default profile
|
|
|
|
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)
|
|
|
|
# Connect to the assigned Neo4j instance
|
|
self.graph = Neo4jGraph(
|
|
url=self.profile.uri,
|
|
username=self.profile.username,
|
|
password=self.profile.password
|
|
)
|
|
|
|
# Load the correct AI model
|
|
self.llm = ChatOpenAI(
|
|
api_key=self.profile.openai_api_key,
|
|
model_name=self.profile.model_name,
|
|
temperature=0
|
|
)
|
|
|
|
# Add the custom prompt to enforce English
|
|
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."
|
|
)
|
|
|
|
# Apply the custom English prompt here
|
|
self.graph_transformer = LLMGraphTransformer(
|
|
llm=self.llm,
|
|
prompt=english_prompt
|
|
)
|
|
|
|
logger.info(f"Neo4jDatabase initialized with profile: {self.profile.name}")
|
|
|
|
except Neo4jProfile.DoesNotExist:
|
|
logger.error(f"Neo4j profile '{profile_name}' not found.")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize Neo4jDatabase: {str(e)}")
|
|
raise
|
|
|
|
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`.
|
|
"""
|
|
try:
|
|
|
|
timestamp = datetime.utcnow().isoformat()
|
|
|
|
# Convert messages into `Document` objects
|
|
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}),
|
|
]
|
|
|
|
# Convert text into structured graph documents
|
|
graph_docs = self.graph_transformer.convert_to_graph_documents(documents)
|
|
|
|
# Store the structured graph data into Neo4j
|
|
self.graph.add_graph_documents(graph_docs, include_source=True)
|
|
|
|
logger.info(f"Stored interaction in Neo4j (Profile: {self.profile.name})")
|
|
except Exception as e:
|
|
logger.error(f"Failed to store interaction in Neo4j: {str(e)}")
|
|
raise
|
|
|
|
def query_graph(self, user_query):
|
|
"""
|
|
Queries the graph using GraphCypherQAChain and returns a structured response.
|
|
"""
|
|
logger.info(f"Calling Neo4j {user_query}")
|
|
try:
|
|
# Use AI model to generate Cypher query
|
|
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 {result}")
|
|
return result['result']
|
|
|
|
except Exception as e:
|
|
logger.error(f"Graph query failed: {e}")
|
|
return None
|