Update chains/arxiv_chains.py

app.py

@@ -22,8 +22,8 @@ from langchain.chains import LLMChain
 from langchain_experimental.utilities.sql_database import SQLDatabase
 from langchain_experimental.retrievers.sql_database import SQLDatabaseChainRetriever
 from langchain_experimental.sql.base import SQLDatabaseChain
-from langchain_experimental.sql.parser import VectorSQLRetrieveAllOutputParser
 
+from chains.arxiv_chains import VectorSQLRetrieveCustomOutputParser
 from chains.arxiv_chains import ArXivQAwithSourcesChain, ArXivStuffDocumentChain
 from callbacks.arxiv_callbacks import ChatDataSelfSearchCallBackHandler, \
     ChatDataSelfAskCallBackHandler, ChatDataSQLSearchCallBackHandler, \
@@ -155,7 +155,7 @@ def build_retriever():
             template=_myscale_prompt,
         )
 
-        output_parser = VectorSQLRetrieveAllOutputParser.from_embeddings(
+        output_parser = VectorSQLRetrieveCustomOutputParser.from_embeddings(
             model=embeddings)
         sql_query_chain = SQLDatabaseChain.from_llm(
             llm=OpenAI(openai_api_key=OPENAI_API_KEY, temperature=0),

chains/arxiv_chains.py

@@ -16,6 +16,23 @@ from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.combine_documents.stuff import StuffDocumentsChain
 
 
+class VectorSQLRetrieveCustomOutputParser(VectorSQLOutputParser):
+    """Based on VectorSQLOutputParser
+    It also modify the SQL to get all columns
+    """
+
+    @property
+    def _type(self) -> str:
+        return "vector_sql_retrieve_custom"
+
+    def parse(self, text: str) -> Dict[str, Any]:
+        text = text.strip()
+        start = text.upper().find("SELECT")
+        if start >= 0:
+            end = text.upper().find("FROM")
+            text = text.replace(text[start + len("SELECT") + 1 : end - 1], "title, abstract, authors, pubdate, categories, id")
+        return super().parse(text)
+
 class ArXivStuffDocumentChain(StuffDocumentsChain):
     """Combine arxiv documents with PDF reference number"""