[email protected] commited on
Commit
89033ee
·
1 Parent(s): a3d26e6

Refactor chain initialization and improve document retrieval handling

Browse files
Files changed (1) hide show
  1. rag.py +7 -6
rag.py CHANGED
@@ -72,19 +72,20 @@ class Rag:
72
  },
73
  )
74
 
75
- self.chain = self.prompt | self.model | StrOutputParser()
76
-
77
  def ask(self, query: str, messages: list):
78
- if not self.chain:
79
- return "Ajouter un document PDF d'abord."
80
 
81
  print("messages ", messages)
82
 
83
  # Retrieve the context document
84
- documentContext = self.retriever.invoke(query)
 
 
 
85
 
86
  # Retrieve the VectoreStore
87
- contextCommon = None
88
 
89
  return self.chain.invoke({
90
  "query": query,
 
72
  },
73
  )
74
 
 
 
75
  def ask(self, query: str, messages: list):
76
+
77
+ self.chain = self.prompt | self.model | StrOutputParser()
78
 
79
  print("messages ", messages)
80
 
81
  # Retrieve the context document
82
+ if self.retriever is None:
83
+ documentContext = ''
84
+ else:
85
+ documentContext = self.retriever.invoke(query)
86
 
87
  # Retrieve the VectoreStore
88
+ contextCommon = self.vector_store.retriever(query, self.embedding)
89
 
90
  return self.chain.invoke({
91
  "query": query,