You can modify which large language model (LLM) you use to suit your own needs.
You may want to:
-
Use a different LLM provider for cost or feature reasons.
-
Adjust model parameters to control the output.
-
Run a local LLM for data privacy or latency reasons.
-
Modify the prompt to cater to your specific domain.
LLM Provider
The LLM you use is configurable in the SimpleKGPipeline via the llm parameter.
The neo4j_graphrag package includes adapters for several popular LLM providers, including OpenAI, Azure OpenAI, and Ollama.
For example, you could run the OpenAI open source model openai/gpt-oss-20b locally using a application such as LM Studio:
llm = OpenAILLM(
model_name="openai/gpt-oss-20b",
model_params={
"temperature": 0
},
base_url = "http://localhost:1234/v1"
)Model Parameters
Throughmodel_params you can change how by model responds by adjusting model parameters such as temperature and response_format.
The parameters available will depend on the specific LLM you are using.You can also create your own LLM adapter by inheriting from the LLMInterface class.
Prompt Customization
The prompt used for entity extraction and other tasks can also be customized by modifying the prompt_template parameter of the SimpleKGPipeline.
You can provide an entirely new prompt, but it is often easier to add to the existing.
If you wanted to restrict entity extraction to a specific domain, such as technology companies, you could modify the prompt as follows:
from neo4j_graphrag.generation.prompts import ERExtractionTemplate
domain_instructions = (
"Only extract entities that are related to the technology industry."
"These include companies, products, programming languages, frameworks, and tools."
"\n"
)
prompt_template = ERExtractionTemplate(
template = domain_instructions + ERExtractionTemplate.DEFAULT_TEMPLATE
)The domain_instructions are added to the start of the default entity extraction prompt to guide the LLM to only extract relevant entities.
The custom prompt can then be used in the SimpleKGPipeline by setting the prompt_template parameter:
kg_builder = SimpleKGPipeline(
llm=llm,
driver=neo4j_driver,
neo4j_database=os.getenv("NEO4J_DATABASE"),
embedder=embedder,
from_pdf=True,
prompt_template=prompt_template,
)Reveal the complete code
This example code shows how to create and use a custom prompt in a SimpleKGPipeline:
import os
from dotenv import load_dotenv
load_dotenv()
import asyncio
from neo4j import GraphDatabase
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.embeddings import OpenAIEmbeddings
from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
from neo4j_graphrag.generation.prompts import ERExtractionTemplate
neo4j_driver = GraphDatabase.driver(
os.getenv("NEO4J_URI"),
auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD"))
)
neo4j_driver.verify_connectivity()
llm = OpenAILLM(
model_name="gpt-4o",
model_params={
"temperature": 0,
"response_format": {"type": "json_object"},
}
)
embedder = OpenAIEmbeddings(
model="text-embedding-ada-002"
)
domain_instructions = (
"Only extract entities that are related to the technology industry."
"These include companies, products, programming languages, frameworks, and tools."
"\n"
)
prompt_template = ERExtractionTemplate(
template = domain_instructions + ERExtractionTemplate.DEFAULT_TEMPLATE
)
kg_builder = SimpleKGPipeline(
llm=llm,
driver=neo4j_driver,
neo4j_database=os.getenv("NEO4J_DATABASE"),
embedder=embedder,
from_pdf=True,
prompt_template=prompt_template,
)
pdf_file = "./genai-graphrag-python/data/genai-fundamentals_1-generative-ai_1-what-is-genai.pdf"
result = asyncio.run(kg_builder.run_async(file_path=pdf_file))
print(result.result)When you’re ready you can continue.
Lesson Summary
In this lesson, you learned about the options for configuring the LLM used in the knowledge graph pipeline, including selecting different LLM providers and customizing prompts.
In the next lesson, you will use what you have learned to create your own knowledge graph from your documents.