Initial commit
This commit is contained in:
@@ -0,0 +1,153 @@
|
||||
from dotenv import load_dotenv, find_dotenv
|
||||
load_dotenv(find_dotenv())
|
||||
import os
|
||||
import modules.logger_tool as logger
|
||||
log_name = 'api_routers_langchain_graph_qa'
|
||||
log_dir = os.getenv("LOG_PATH", "/logs") # Default path as fallback
|
||||
logging = logger.get_logger(
|
||||
name=log_name,
|
||||
log_level=os.getenv("LOG_LEVEL", "DEBUG"),
|
||||
log_path=log_dir,
|
||||
log_file=log_name,
|
||||
runtime=True,
|
||||
log_format='default'
|
||||
)
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from langchain.chains import GraphCypherQAChain
|
||||
from langchain_community.graphs import Neo4jGraph
|
||||
from langchain_community.chat_models import ChatOpenAI
|
||||
from langchain.prompts.prompt import PromptTemplate
|
||||
from routers.llm.private.ollama.ollama_wrapper import OllamaWrapper
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# Define the schema for nodes and relationships
|
||||
node_types = {
|
||||
"KeyStage": ["merged", "key_stage_name", "unique_id", "created"],
|
||||
"KeyStageSyllabus": ["ks_syllabus_name", "unique_id", "created", "merged", "ks_syllabus_key_stage", "ks_syllabus_subject"],
|
||||
"YearGroup": ["created", "merged", "unique_id", "year_group_name"],
|
||||
"YearGroupSyllabus": ["created", "merged", "yr_syllabus_name", "yr_syllabus_year_group", "yr_syllabus_id", "yr_syllabus_subject"],
|
||||
"Topic": ["topic_type", "topic_assessment_type", "created", "merged", "unique_id", "topic_id", "total_number_of_lessons_for_topic", "topic_title"],
|
||||
"Lesson": ["topic_lesson_id", "topic_lesson_type", "created", "merged", "topic_lesson_title", "topic_lesson_length", "topic_lesson_suggested_activities", "topic_lesson_weblinks", "topic_lesson_skills_learned"],
|
||||
"LearningStatement": ["created", "merged", "lesson_learning_statement", "lesson_learning_statement_id", "lesson_learning_statement_type"]
|
||||
}
|
||||
|
||||
relationship_types = {
|
||||
"KEY_STAGE_INCLUDES_KEY_STAGE_SYLLABUS": ["created", "merged"],
|
||||
"KEY_STAGE_SYLLABUS_INCLUDES_YEAR_GROUP_SYLLABUS": ["created", "merged"],
|
||||
"YEAR_GROUP_FOLLOWS_YEAR_GROUP": ["created", "merged"],
|
||||
"KEY_STAGE_FOLLOWS_KEY_STAGE": ["created", "merged"],
|
||||
"YEAR_SYLLABUS_INCLUDES_TOPIC": ["created", "merged"],
|
||||
"TOPIC_INCLUDES_LESSON": ["created", "merged"],
|
||||
"LESSON_INCLUDES_LEARNING_STATEMENT": ["created", "merged"],
|
||||
"LESSON_FOLLOWS_LESSON": ["created", "merged"]
|
||||
}
|
||||
|
||||
@router.get("/prompt")
|
||||
async def query_graph(
|
||||
database: str, prompt: str, top_k: int = 30, model: str = "gpt-4o", temperature: float = 0,
|
||||
verbose: bool = False, return_intermediate_steps: bool = False, exclude_types: list = None, include_types: list = None,
|
||||
return_direct: bool = False, validate_cypher: bool = False, model_type: str = "openai"
|
||||
):
|
||||
logging.info(f"Received request with prompt: {prompt}")
|
||||
if exclude_types is None:
|
||||
logging.info("No exclude_types provided, using default.")
|
||||
exclude_types = []
|
||||
if include_types is None:
|
||||
logging.info("No include_types provided, using default.")
|
||||
include_types = []
|
||||
|
||||
# Validate include_types and exclude_types
|
||||
logging.info(f"Validating include_types and exclude_types...")
|
||||
valid_types = set(node_types.keys()).union(set(relationship_types.keys()))
|
||||
logging.info(f"Valid types: {valid_types}")
|
||||
exclude_types = [t for t in exclude_types if t in valid_types]
|
||||
logging.info(f"Validated exclude_types: {exclude_types}")
|
||||
include_types = [t for t in include_types if t in valid_types]
|
||||
logging.info(f"Validated include_types: {include_types}")
|
||||
|
||||
graph = Neo4jGraph(
|
||||
url=os.environ['APP_BOLT_URL'],
|
||||
username=os.environ['USER_NEO4J'],
|
||||
password=os.environ['PASSWORD_NEO4J'],
|
||||
database=database
|
||||
)
|
||||
|
||||
logging.info("Refreshing schema...")
|
||||
graph.refresh_schema()
|
||||
logging.info("Schema refreshed.")
|
||||
schema = graph.schema
|
||||
logging.info(f"Schema: {schema}")
|
||||
|
||||
CYPHER_GENERATION_TEMPLATE = """Task: Generate a Cypher statement to query a graph database for timetable information.
|
||||
Role:
|
||||
You are an assistant in a school for teachers, specializing in querying graph databases to find answers to questions.
|
||||
The teacher will ask you questions about their timetable.
|
||||
|
||||
Instructions:
|
||||
1. Use only the provided relationship types and properties in the schema.
|
||||
2. Do not use any other relationship types or properties that are not provided.
|
||||
|
||||
Schema:
|
||||
{schema}
|
||||
|
||||
Note:
|
||||
1. Do not include any explanations or apologies in your responses.
|
||||
2. Do not respond to any questions that might ask anything else than for you to construct a Cypher statement.
|
||||
3. Do not include any text except the generated Cypher statement.
|
||||
|
||||
The question is:
|
||||
{question}"""
|
||||
|
||||
CYPHER_GENERATION_PROMPT = PromptTemplate(
|
||||
input_variables=["schema", "question"],
|
||||
template=CYPHER_GENERATION_TEMPLATE
|
||||
)
|
||||
|
||||
if model_type == "ollama":
|
||||
ollama_host = os.getenv("OLLAMA_URL")
|
||||
ollama_port = os.getenv("OLLAMA_PORT")
|
||||
if not ollama_host or not ollama_port:
|
||||
raise HTTPException(status_code=500, detail="Ollama host or port not set")
|
||||
client = OllamaWrapper(host=f'http://{ollama_host}:{ollama_port}')
|
||||
cypher_llm = client
|
||||
qa_llm = client
|
||||
else:
|
||||
cypher_llm = ChatOpenAI(temperature=temperature, model=model)
|
||||
qa_llm = ChatOpenAI(temperature=temperature, model=model)
|
||||
|
||||
chain = GraphCypherQAChain.from_llm(
|
||||
graph=graph,
|
||||
cypher_llm=cypher_llm,
|
||||
qa_llm=qa_llm,
|
||||
top_k=top_k,
|
||||
verbose=verbose,
|
||||
cypher_prompt=CYPHER_GENERATION_PROMPT,
|
||||
return_intermediate_steps=return_intermediate_steps,
|
||||
exclude_types=exclude_types,
|
||||
include_types=include_types,
|
||||
return_direct=return_direct,
|
||||
validate_cypher=validate_cypher
|
||||
)
|
||||
|
||||
formatted_prompt = CYPHER_GENERATION_PROMPT.format(schema=schema, question=prompt)
|
||||
|
||||
logging.info("\n\n")
|
||||
|
||||
logging.info("==================================================")
|
||||
logging.info("= graph_qa.py =")
|
||||
logging.info("==================================================")
|
||||
logging.info(f"Prompt: {prompt}")
|
||||
logging.info("--------------------------------------------------")
|
||||
logging.info(f"Schema: \n{schema}\n")
|
||||
logging.info("--------------------------------------------------")
|
||||
logging.info(f"Formatted Prompt: \n{formatted_prompt}\n")
|
||||
logging.info("--------------------------------------------------")
|
||||
logging.info(f"Cypher prompt: \n{CYPHER_GENERATION_PROMPT}\n")
|
||||
logging.info("--------------------------------------------------")
|
||||
logging.info(f"Cypher template: \n{CYPHER_GENERATION_TEMPLATE}\n")
|
||||
logging.info("--------------------------------------------------")
|
||||
logging.info(f"Cypher chain: \n{chain}\n")
|
||||
logging.info("==================================================")
|
||||
|
||||
return chain(prompt)
|
||||
Reference in New Issue
Block a user