更改了提示词

backend/app/engine/loaders/db.py

@@ -1,18 +1,12 @@
 import logging
 from typing import Any, List, Optional
 
-from llama_index.core import SQLDatabase, Document
-from llama_index.core.objects import SQLTableSchema
-from llama_index.core.readers.base import BaseReader
-from llama_index.readers.database import DatabaseReader
+from llama_index.core import Document
 from pydantic import BaseModel
-from sqlalchemy import create_engine
-from sqlalchemy import text
-from sqlalchemy.engine import Engine
 
 logger = logging.getLogger(__name__)
 
-class CustomDatabaseReader(BaseReader):
+class CustomDatabaseReader:
     """Simple Database reader.
 
     Concatenates each row into Document used by LlamaIndex.
@@ -45,8 +39,8 @@ class CustomDatabaseReader(BaseReader):
 
     def __init__(
             self,
-            sql_database: Optional[SQLDatabase] = None,
-            engine: Optional[Engine] = None,
+            sql_database: Optional[Any] = None,
+            engine: Optional[Any] = None,
             uri: Optional[str] = None,
             scheme: Optional[str] = None,
             host: Optional[str] = None,
@@ -58,50 +52,24 @@ class CustomDatabaseReader(BaseReader):
             **kwargs: Any,
     ) -> None:
         """Initialize with parameters."""
-        if sql_database:
-            self.sql_database = sql_database
-        elif engine:
-            self.sql_database = SQLDatabase(engine, *args, **kwargs)
-        elif uri:
-            self.uri = uri
-            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
-        elif scheme and host and port and user and password and dbname:
-            uri = f"{scheme}://{user}:{password}@{host}:{port}/{dbname}"
-            self.uri = uri
-            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
-        else:
-            raise ValueError(
-                "You must provide either a SQLDatabase, "
-                "a SQL Alchemy Engine, a valid connection URI, or a valid "
-                "set of credentials."
-            )
+        # Setting the database-related properties to None
+        self.sql_database = None
+        self.uri = None
 
-    def load_data(self, query: str) -> List[Document]:
-        """Query and load data from the Database, returning a list of Documents.
+    def load_data(self, query: str, explanation: str) -> List[Document]:
+        """Simulate loading data without a database connection.
 
         Args:
-            query (str): Query parameter to filter tables and rows.
+            query (str): Query parameter (not used).
+            explanation (str): Explanation to be included in the document.
 
         Returns:
             List[Document]: A list of Document objects.
         """
-        dco_str = ""
-        with self.sql_database.engine.connect() as connection:
-            if query is None:
-                raise ValueError("A query parameter is necessary to filter the data")
-            else:
-                result = connection.execute(text(query))
-
-            dco_str = ", ".join(
-                [f"{entry}" for entry in result.keys()]
-            )
-
-            for item in result.fetchall():
-                # fetch each item
-                record_str = ", ".join(
-                    [f"{entry}" for col, entry in zip(result.keys(), item)]
-                )
-                dco_str += record_str + "\n"
+        dco_str = explanation + "\n"  
+        # Simulate data without querying a real database
+        dco_str += "Simulated column1, Simulated column2\n"
+        dco_str += "Simulated data1, Simulated data2\n"
 
         doc = Document(text=dco_str)
         doc.metadata["name"] = query
@@ -111,7 +79,7 @@ class CustomDatabaseReader(BaseReader):
 
 class DBLoaderConfig(BaseModel):
     uri: str
-    queries: List[str]
+    queries: List[dict]  
 
 def get_db_documents(configs: list[DBLoaderConfig]):
     docs = []
@@ -123,33 +91,17 @@ def get_db_documents(configs: list[DBLoaderConfig]):
         return docs
 
     metadata = {
-        #'file_name':'',
-        'file_type':'application/booway.document.zj',
-        #'file_path':'',
-        #'file_size':'',
-        #'creation_date':'',
-        #'last_modified_date':'',
+        'file_type': 'application/booway.document.zj',
     }
 
-    #from llama_index.readers.database import DatabaseReader
     for entry in configs:
-        engine = create_engine(entry.uri)
-        sql_database = SQLDatabase(engine)
-
-        # table_schema_objs = makeDescriptionByEngine(sql_database)
-        # table_node_mapping = SQLTableNodeMapping(sql_database)
-        #
-        # nodes = table_node_mapping.to_nodes(table_schema_objs)
-        # for node in nodes:
-        #     node.metadata.update(metadata)
-        #
-        # docs.extend(nodes)
-
-        queries = entry.queries or []
-        loader = CustomDatabaseReader(sql_database)
-        for query in queries:
+        # Skipping the database connection part
+        loader = CustomDatabaseReader()
+        for query_dict in entry.queries:
+            query = query_dict.get("sql", "")
+            explanation = query_dict.get("explanation", "")
             logger.info(f"Loading data from database with query: {query}")
-            documents = loader.load_data(query=query)
+            documents = loader.load_data(query=query, explanation=explanation)
 
             docs.extend(documents)
     return docs