src / tools.py
msz-archive's picture
Upload folder using huggingface_hub
28295d6 verified
from typing import List, Dict, Optional
from langchain.tools import StructuredTool
import json
# Now import from payments module
from src.payments.agent import process_purchase_request
# Load datasets
with open("data/supplements.json", "r") as f:
SUPPLEMENTS = json.load(f)
with open("data/prescription_drugs.json", "r") as f:
PRESCRIPTION_DRUGS = json.load(f)
# Combine both datasets for unified searching
ALL_PRODUCTS = {**SUPPLEMENTS, **PRESCRIPTION_DRUGS}
# Supplement search and retrieval functions
def search_supplements_by_condition(query: str) -> List[Dict]:
"""
Search for supplements that are recommended for a specific health condition,
well-being goal, or enhancement purpose.
Args:
query: The condition, goal, or purpose to search supplements for
(e.g., 'hair_loss', 'cognitive_function', 'muscle_growth')
Returns:
A list of supplements that match the query
"""
matching_supplements = []
query = query.lower()
for supp_id, supp_data in SUPPLEMENTS.items():
# Check conditions
conditions = [c.lower() for c in supp_data.get("conditions", [])]
# Check benefits/enhancements
benefits = [b.lower() for b in supp_data.get("benefits", [])]
# Check categories
categories = [cat.lower() for cat in supp_data.get("categories", [])]
# Match against all fields
if (query in conditions or
query in benefits or
query in categories or
query in supp_data["name"].lower() or
query in supp_data["description"].lower()):
matching_supplements.append({
"id": supp_id,
"name": supp_data["name"],
"description": supp_data["description"],
"link": supp_data["affiliate_link"],
"conditions": supp_data.get("conditions", []),
"benefits": supp_data.get("benefits", []),
"categories": supp_data.get("categories", []),
"disclaimers": supp_data.get("disclaimers", [])
})
return matching_supplements
def search_supplements_by_name(name: str) -> Optional[Dict]:
"""
Search for a supplement by its name.
Args:
name: The name of the supplement to search for (e.g., 'Vitamin D3', 'Magnesium')
Returns:
Details of the supplement if found, None otherwise
"""
name = name.lower()
for supp_id, supp_data in SUPPLEMENTS.items():
if name in supp_data["name"].lower():
return {
"id": supp_id,
"name": supp_data["name"],
"description": supp_data["description"],
"link": supp_data["affiliate_link"],
"conditions": supp_data.get("conditions", []),
"benefits": supp_data.get("benefits", []),
"categories": supp_data.get("categories", []),
"disclaimers": supp_data.get("disclaimers", [])
}
return None
def get_supplement_details(supplement_id: str) -> Optional[Dict]:
"""
Get detailed information about a specific supplement.
Args:
supplement_id: The ID of the supplement to retrieve
Returns:
Detailed information about the supplement or None if not found
"""
if supplement_id in SUPPLEMENTS:
supp_data = SUPPLEMENTS[supplement_id]
return {
"id": supplement_id,
"name": supp_data["name"],
"description": supp_data["description"],
"link": supp_data["affiliate_link"],
"conditions": supp_data.get("conditions", []),
"disclaimers": supp_data.get("disclaimers", [])
}
return None
def list_all_supplements() -> List[Dict]:
"""
List all available supplements in the database.
Returns:
A list of all supplements with their basic information
"""
return [
{
"id": supp_id,
"name": supp_data["name"],
"description": supp_data["description"],
"conditions": supp_data.get("conditions", [])
}
for supp_id, supp_data in SUPPLEMENTS.items()
]
# Prescription drug search and retrieval functions
def search_prescription_drugs_by_condition(query: str) -> List[Dict]:
"""
Search for prescription drugs that are used for a specific health condition,
well-being goal, or therapeutic purpose.
Args:
query: The condition, goal, or purpose to search drugs for
(e.g., 'hypertension', 'depression', 'hair_loss')
Returns:
A list of prescription drugs that match the query
"""
matching_drugs = []
query = query.lower()
for drug_id, drug_data in PRESCRIPTION_DRUGS.items():
# Check conditions
conditions = [c.lower() for c in drug_data.get("conditions", [])]
# Check benefits
benefits = [b.lower() for b in drug_data.get("benefits", [])]
# Check categories
categories = [cat.lower() for cat in drug_data.get("categories", [])]
# Match against all fields
if (query in conditions or
query in benefits or
query in categories or
query in drug_data["name"].lower() or
query in drug_data["description"].lower()):
matching_drugs.append({
"id": drug_id,
"name": drug_data["name"],
"description": drug_data["description"],
"link": drug_data["affiliate_link"],
"conditions": drug_data.get("conditions", []),
"benefits": drug_data.get("benefits", []),
"categories": drug_data.get("categories", []),
"disclaimers": drug_data.get("disclaimers", []),
"requires_prescription": True
})
return matching_drugs
def search_prescription_drugs_by_name(name: str) -> Optional[Dict]:
"""
Search for a prescription drug by its name.
Args:
name: The name of the drug to search for (e.g., 'Metformin', 'Atorvastatin')
Returns:
Details of the prescription drug if found, None otherwise
"""
name = name.lower()
for drug_id, drug_data in PRESCRIPTION_DRUGS.items():
if name in drug_data["name"].lower():
return {
"id": drug_id,
"name": drug_data["name"],
"description": drug_data["description"],
"link": drug_data["affiliate_link"],
"conditions": drug_data.get("conditions", []),
"benefits": drug_data.get("benefits", []),
"categories": drug_data.get("categories", []),
"disclaimers": drug_data.get("disclaimers", []),
"requires_prescription": True
}
return None
def get_prescription_drug_details(drug_id: str) -> Optional[Dict]:
"""
Get detailed information about a specific prescription drug.
Args:
drug_id: The ID of the prescription drug to retrieve
Returns:
Detailed information about the drug or None if not found
"""
if drug_id in PRESCRIPTION_DRUGS:
drug_data = PRESCRIPTION_DRUGS[drug_id]
return {
"id": drug_id,
"name": drug_data["name"],
"description": drug_data["description"],
"link": drug_data["affiliate_link"],
"conditions": drug_data.get("conditions", []),
"benefits": drug_data.get("benefits", []),
"categories": drug_data.get("categories", []),
"disclaimers": drug_data.get("disclaimers", []),
"requires_prescription": True
}
return None
def list_all_prescription_drugs() -> List[Dict]:
"""
List all available prescription drugs in the database.
Returns:
A list of all prescription drugs with their basic information
"""
return [
{
"id": drug_id,
"name": drug_data["name"],
"description": drug_data["description"],
"conditions": drug_data.get("conditions", []),
"requires_prescription": True
}
for drug_id, drug_data in PRESCRIPTION_DRUGS.items()
]
# Combined search functions
def search_all_products_by_condition(query: str) -> List[Dict]:
"""
Search for both supplements and prescription drugs that are recommended for a specific health condition,
well-being goal, or enhancement purpose.
Args:
query: The condition, goal, or purpose to search for
(e.g., 'hair_loss', 'cognitive_function', 'depression')
Returns:
A list of supplements and prescription drugs that match the query
"""
supplements = search_supplements_by_condition(query)
prescription_drugs = search_prescription_drugs_by_condition(query)
return supplements + prescription_drugs
def search_all_products_by_name(name: str) -> List[Dict]:
"""
Search for both supplements and prescription drugs by name.
Args:
name: The name to search for (e.g., 'Vitamin D3', 'Metformin')
Returns:
List of products matching the name query
"""
results = []
supplement = search_supplements_by_name(name)
if supplement:
results.append(supplement)
prescription = search_prescription_drugs_by_name(name)
if prescription:
results.append(prescription)
return results
def get_product_details(product_id: str) -> Optional[Dict]:
"""
Get detailed information about a specific product (supplement or prescription drug).
Args:
product_id: The ID of the product to retrieve
Returns:
Detailed information about the product or None if not found
"""
supplement = get_supplement_details(product_id)
if supplement:
return supplement
prescription = get_prescription_drug_details(product_id)
if prescription:
return prescription
return None
def list_all_products() -> List[Dict]:
"""
List all available products (supplements and prescription drugs) in the database.
Returns:
A list of all products with their basic information
"""
supplements = list_all_supplements()
prescription_drugs = list_all_prescription_drugs()
return supplements + prescription_drugs
def create_payment_link(request: str) -> str:
"""
Create a payment link for products and consultations based on a natural language request.
Args:
request: A string describing what products and quantities to purchase
(e.g., "I want to buy 1 tretinoin and 1 consultation")
Returns:
The payment link URL
"""
# Process the request and get payment link
response = process_purchase_request(request)
#response['messages'][0].content
return response
# Create and export Langchain tools
def get_tools():
"""Return all the tools needed for the agent"""
return [
StructuredTool.from_function(search_supplements_by_condition),
StructuredTool.from_function(search_prescription_drugs_by_condition),
StructuredTool.from_function(search_all_products_by_condition),
StructuredTool.from_function(search_supplements_by_name),
StructuredTool.from_function(search_prescription_drugs_by_name),
StructuredTool.from_function(search_all_products_by_name),
StructuredTool.from_function(get_supplement_details),
StructuredTool.from_function(get_prescription_drug_details),
StructuredTool.from_function(get_product_details),
StructuredTool.from_function(list_all_supplements),
StructuredTool.from_function(list_all_prescription_drugs),
StructuredTool.from_function(list_all_products),
StructuredTool.from_function(create_payment_link)
]
# Export data for use in main.py
def get_supplements():
return SUPPLEMENTS
def get_prescription_drugs():
return PRESCRIPTION_DRUGS
def get_available_products() -> Dict:
"""
Get a list of all available products in the catalog.
Returns:
A dictionary containing available products categorized by type
"""
# Load the products catalog
with open("src/payments/products.json", "r") as f:
products_catalog = json.load(f)
# Separate into supplements, prescription drugs, and services
supplements = []
prescription_drugs = []
services = []
for product_id in products_catalog.keys():
if product_id == "NP_consultation":
services.append(product_id) # Add consultation as a service
elif product_id in PRESCRIPTION_DRUGS:
prescription_drugs.append(product_id)
elif product_id in SUPPLEMENTS:
supplements.append(product_id)
return {
"supplements": supplements,
"prescription_drugs": prescription_drugs,
"services": services,
"all_products": list(products_catalog.keys())
}
if __name__ == "__main__":
sample_request = '1 Finasteride consultation, 1 Minoxidil consultation'
result = process_purchase_request(sample_request)
print(result)
breakpoint()
asd = 1
# You can test with other examples
# result2 = process_purchase_request("I want to buy 3 items of product_i and 2 of product_k")
# print(result2)