morphik-core/examples/prompt_overrides.py
Adityavardhan Agrawal 1792275cb8
Format fix, UI package update (#100)
Co-authored-by: Arnav Agrawal <aa779@cornell.edu>
2025-04-20 16:34:29 -07:00

117 lines
4.3 KiB
Python

import os
from dotenv import load_dotenv
from morphik import Morphik
from morphik.models import (
EntityExtractionExample,
EntityExtractionPromptOverride,
EntityResolutionExample,
EntityResolutionPromptOverride,
GraphPromptOverrides,
QueryPromptOverride,
QueryPromptOverrides,
)
# Load environment variables
load_dotenv()
# Connect to Morphik
db = Morphik(os.getenv("MORPHIK_URI"), timeout=10000, is_local=True)
# Ingest some sample medical documents
medical_texts = [
{
"text": "Patients with Type 2 Diabetes often show increased insulin resistance. Treatment options include metformin and lifestyle changes.",
"metadata": {"category": "medical", "specialty": "endocrinology"},
},
{
"text": "Hypertension (high blood pressure) is a common comorbidity of diabetes. ACE inhibitors are frequently prescribed to manage blood pressure.",
"metadata": {"category": "medical", "specialty": "cardiology"},
},
{
"text": "Studies show that regular exercise can improve glucose control in diabetic patients and reduce the risk of cardiovascular disease.",
"metadata": {"category": "medical", "specialty": "research"},
},
]
# Ingest documents
doc_ids = []
for item in medical_texts:
doc = db.ingest_text(item["text"], metadata=item["metadata"])
doc_ids.append(doc.external_id)
print(f"Ingested document with ID: {doc.external_id}")
# Example 1: Basic Query with Prompt Override
print("\nExample 1: Basic Query with Prompt Override")
basic_query_override = QueryPromptOverrides(
query=QueryPromptOverride(
prompt_template="Respond as if you are a medical professional. Answer the following question based on the provided context: {question}"
)
)
response = db.query("What treatments are available for diabetes?", prompt_overrides=basic_query_override)
print("Response with medical professional prompt:")
print(response.completion)
# Example 2: Query without override for comparison
print("\nExample 2: Same Query without Override")
standard_response = db.query("What treatments are available for diabetes?")
print("Standard response:")
print(standard_response.completion)
# Example 3: Create a knowledge graph with customized entity extraction
print("\nExample 3: Knowledge Graph with Customized Entity Extraction")
graph_overrides = GraphPromptOverrides(
entity_extraction=EntityExtractionPromptOverride(
examples=[
EntityExtractionExample(label="Diabetes", type="CONDITION"),
EntityExtractionExample(label="Metformin", type="MEDICATION"),
EntityExtractionExample(label="Hypertension", type="CONDITION"),
EntityExtractionExample(label="ACE inhibitors", type="MEDICATION"),
]
),
entity_resolution=EntityResolutionPromptOverride(
examples=[
EntityResolutionExample(canonical="Diabetes Mellitus", variants=["Diabetes", "Type 2 Diabetes", "T2DM"]),
EntityResolutionExample(canonical="Hypertension", variants=["High Blood Pressure", "HTN", "Elevated BP"]),
]
),
)
graph = db.create_graph(name="medical_conditions_graph", documents=doc_ids, prompt_overrides=graph_overrides)
print(f"Created graph with {len(graph.entities)} entities and {len(graph.relationships)} relationships")
# Example 4: Query using the customized graph
print("\nExample 4: Query Using the Customized Graph")
graph_response = db.query(
"How are diabetes and hypertension related?",
graph_name="medical_conditions_graph",
hop_depth=2,
include_paths=True,
)
print("Response using knowledge graph:")
print(graph_response.completion)
# Print the relationship paths if available
if graph_response.metadata and "graph" in graph_response.metadata:
print("\nRelationship paths:")
for path in graph_response.metadata["graph"]["paths"]:
print(" -> ".join(path))
# Example 5: Using dictionary for prompt overrides
print("\nExample 5: Using Dictionary for Prompt Overrides")
dict_override = {
"query": {"prompt_template": "Summarize the information in a bulleted list format, focusing on: {question}"}
}
dict_response = db.query(
"What are the connections between diabetes, medications, and exercise?",
prompt_overrides=dict_override,
)
print("Response with dictionary-based prompt override:")
print(dict_response.completion)