Spaces:
Sleeping
Sleeping
import os | |
import streamlit as st | |
from langchain_community.graphs import Neo4jGraph | |
from langchain_groq import ChatGroq | |
from langchain.chains import GraphCypherQAChain | |
from dotenv import load_dotenv | |
# Load environment variables | |
load_dotenv() | |
# Load environment variables | |
NEO4J_URI = "neo4j+s://64beefae.databases.neo4j.io" | |
NEO4J_USERNAME = "neo4j" | |
NEO4J_PASSWORD = "OTj5yGnWLF59yx4UX1g3xABarVOvVKiM3CT9L4bNkF8" | |
GROQ_API_KEY = "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" | |
# Initialize Neo4j graph | |
graph = Neo4jGraph(url=NEO4J_URI, username=NEO4J_USERNAME, password=NEO4J_PASSWORD) | |
# Load movie dataset into Neo4j | |
moview_query = """ | |
LOAD CSV WITH HEADERS FROM | |
'https://raw.githubusercontent.com/tomasonjo/blog-datasets/main/movies/movies_small.csv' as row | |
MERGE(m:Movie{id:row.movieId}) | |
SET m.released = date(row.released), | |
m.title = row.title, | |
m.imdbRating = toFloat(row.imdbRating) | |
FOREACH (director in split(row.director, '|') | | |
MERGE (p:Person {name:trim(director)}) | |
MERGE (p)-[:DIRECTED]->(m)) | |
FOREACH (actor in split(row.actors, '|') | | |
MERGE (p:Person {name:trim(actor)}) | |
MERGE (p)-[:ACTED_IN]->(m)) | |
FOREACH (genre in split(row.genres, '|') | | |
MERGE (g:Genre {name:trim(genre)}) | |
MERGE (m)-[:IN_GENRE]->(g)) | |
""" | |
graph.query(moview_query) | |
graph.refresh_schema() | |
# Initialize LLM | |
llm = ChatGroq(groq_api_key=GROQ_API_KEY, model_name="Gemma2-9b-It") | |
# Create QA Chain | |
chain = GraphCypherQAChain.from_llm(graph=graph, llm=llm, verbose=True) | |
# Streamlit UI | |
st.title("Movie Director Finder") | |
query = st.text_input("Enter your query:", "Who was the director of the movie Casino") | |
if st.button("Find Director"): | |
response = chain.invoke({"query": query}) | |
# Debug: Print the raw response to check its structure | |
st.write("Raw Response:", response) | |
# Extract the director's name from the full context | |
if response and 'full_context' in response: | |
full_context = response['full_context'] | |
if full_context and isinstance(full_context, list) and len(full_context) > 0: | |
director_info = full_context[0] | |
director_name = director_info.get('p.name', None) | |
if director_name: | |
st.write(f"The director of the movie is {director_name}.") | |
else: | |
st.write("Could not find the director's name in the response.") | |
else: | |
st.write("No full context found in the response.") | |
else: | |
st.write("I don't know the answer.") | |