50 lines
1.2 KiB
Python
50 lines
1.2 KiB
Python
"""The LLM interface is compatible with LangChain chat API,
|
|
so any LangChain implementation can be used. Same for embedders.
|
|
For instance, in GraphRAG:
|
|
|
|
Requires OPENAI_API_KEY to be in the env var.
|
|
"""
|
|
|
|
import neo4j
|
|
from langchain_openai.chat_models import ChatOpenAI
|
|
from langchain_openai.embeddings import OpenAIEmbeddings
|
|
from neo4j_graphrag.generation import GraphRAG
|
|
from neo4j_graphrag.retrievers import VectorCypherRetriever
|
|
|
|
# Define database credentials
|
|
URI = "neo4j+s://demo.neo4jlabs.com"
|
|
AUTH = ("recommendations", "recommendations")
|
|
DATABASE = "recommendations"
|
|
INDEX = "moviePlotsEmbedding"
|
|
|
|
|
|
driver = neo4j.GraphDatabase.driver(
|
|
URI,
|
|
auth=AUTH,
|
|
)
|
|
|
|
embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
|
|
|
|
retriever = VectorCypherRetriever(
|
|
driver,
|
|
index_name=INDEX,
|
|
retrieval_query="WITH node, score RETURN node.title as title, node.plot as plot",
|
|
embedder=embedder, # type: ignore[arg-type, unused-ignore]
|
|
neo4j_database=DATABASE,
|
|
)
|
|
|
|
llm = ChatOpenAI(model="gpt-5", temperature=0)
|
|
|
|
rag = GraphRAG(
|
|
retriever=retriever,
|
|
llm=llm, # type: ignore[arg-type, unused-ignore]
|
|
)
|
|
|
|
result = rag.search(
|
|
"Tell me more about Avatar movies",
|
|
return_context=False,
|
|
)
|
|
print(result.answer)
|
|
|
|
driver.close()
|