22 lines
718 B
Python
22 lines
718 B
Python
from langchain_core.documents import Document
|
|
from langchain_core.language_models import BaseChatModel
|
|
|
|
|
|
def categorize(doc: Document, llm: BaseChatModel) -> set[str]:
|
|
# Create a prompt for category extraction
|
|
prompt = f"""
|
|
Extract up to 3 relevant categories from the following document.
|
|
Return only the category names as a list of JSON strings.
|
|
|
|
If you cannot find any relevant categories, return an empty list.
|
|
|
|
Title: {doc.metadata.get('title', 'No title')}
|
|
Content: {doc.page_content}...
|
|
|
|
Categories:"""
|
|
|
|
# Get response from LLM
|
|
result = llm.with_structured_output(method="json_mode").invoke(prompt)
|
|
categories = result.get("categories", [])
|
|
return set(categories)
|