From 7642579f2920682ec33e6775dc66a9c801cfaff3 Mon Sep 17 00:00:00 2001 From: KevKibe Date: Tue, 30 Apr 2024 19:22:13 +0300 Subject: [PATCH] add sources, titles and page_content to QueryResult Model --- src/utils/response_model.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/src/utils/response_model.py b/src/utils/response_model.py index fee8440..3706e20 100644 --- a/src/utils/response_model.py +++ b/src/utils/response_model.py @@ -13,7 +13,28 @@ class QueryResult(BaseModel): page: Union[float, int] = Field(..., description="The page number of the final result of the query.") source_documents: List[Document] = Field(..., description="A list of source documents related to the query.") -# Example usage: + @property + def sources(self) -> List[Union[float, int]]: + """ + Returns a list of the sources (page numbers) from the source documents. + """ + return [doc.source for doc in self.source_documents] + + @property + def titles(self) -> List[str]: + """ + Returns a list of the titles from the source documents. + """ + return [doc.title for doc in self.source_documents] + + @property + def page_contents(self) -> List[str]: + """ + Returns a list of the page contents from the source documents. + """ + return [doc.page_content for doc in self.source_documents] + +# Example # data = { # 'query': 'how did RAG come up?', # 'result': 'RAG came up as a language model that is more strongly grounded in '