File size: 1,334 Bytes
a8ab4cd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import config
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
## Grading model
class DocumentGrade(BaseModel):
"""Binary score for relevance check on retrieved documents"""
binary_score: str = Field(
description='Document is relevant to the question, "yes" or "no"'
)
# Grader Prompts
system = """You are a grader assessing relevance of a retrieved document to a user question. \n
If the document contains keyword(s) or semantic meaning related to the question, grade it as relevant. \n
Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""
grade_prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
]
)
# LLM with function call
llm = ChatOpenAI(model=config.GRADER_MODEL, temperature=0, streaming=True)
structured_llm_grader = llm.with_structured_output(DocumentGrade)
structured_llm_grader
retrieval_grader = grade_prompt | structured_llm_grader
# question = "agent memory"
# docs = retriever.invoke(question)
# doc_txt = docs[1].page_content
# print(retrieval_grader.invoke({"question": question, "document": doc_txt}))
|