from llama_index.core.query_pipeline import ( QueryPipeline, Link, InputComponent, CustomQueryComponent, ) from llama_index.core.objects import ( SQLTableNodeMapping, ObjectIndex, SQLTableSchema, ) from pyvis.network import Network import Stemmer from IPython.display import display, HTML from sqlalchemy import create_engine from llama_index.core import SQLDatabase, VectorStoreIndex, PromptTemplate from llama_index.core.program import LLMTextCompletionProgram from llama_index.core.bridge.pydantic import BaseModel, Field from typing import Dict, List, Any from llama_index.core.query_pipeline import FnComponent from llama_index.core.prompts.default_prompts import DEFAULT_TEXT_TO_SQL_PROMPT from llama_index.core.retrievers import SQLRetriever from llama_index.llms.ollama import Ollama from llama_index.core.objects.base import ObjectRetriever import pymysql, pandas as pd from llama_index.retrievers.bm25 import BM25Retriever from llama_index.core.schema import IndexNode from modules import ( get_table_obj_retriever, create_table_obj_retriever, get_table_context_str, parse_response_to_sql, CustomSQLRetriever ) PROMPT_STR = """\ Give me a summary of the table with the following format. - table_summary: Describe what the table is about in short. Columns: [col1(type), col2(type), ...] Table: {table_str} """ db_user = "shenzhen_ai_for_vibemate_eson" db_password = "dBsnc7OrM0MVi0FEhiHe2y" db_host = "192.168.1.99" db_port = 3306 db_name = "hytto_surfease" TABLE_SUMMARY = { "t_sur_media_sync_es": "This table is about Porn video information:\n\nt_sur_media_sync_es: Columns:id (integer), web_url (string), duration (integer), pattern_per (integer), like_count (integer), dislike_count (integer), view_count (integer), cover_picture (string), title (string), upload_date (datetime), uploader (string), create_time (datetime), update_time (datetime), categories (list of strings), abbreviate_video_url (string), abbreviate_mp4_video_url (string), resource_type (string), like_count_show (integer), stat_version (integer), tags (list of strings), model_name (string), publisher_type (string), period (integer), sexual_preference (string), country (string), type (string), rank_number (integer), rank_rate (float), has_pattern (boolean), trace (string), manifest_url (string), is_delete (boolean), web_url_md5 (string), view_key (string)", "t_sur_models_info": "This table is about Stripchat models' information:\n\nt_sur_models_info: Columns:id (INTEGER), username (VARCHAR(100), image (VARCHAR(500), num_users (INTEGER), pf (VARCHAR(50), pf_model_unite (VARCHAR(50), use_plugin (INTEGER), create_time (DATETIME), update_time (DATETIME), update_time (DATETIME), gender (VARCHAR(50), broadcast_type (VARCHAR(50), common_gender (VARCHAR(50), avatar (VARCHAR(512), age (INTEGER) " } class SQLPipeline: def __init__(self, llm: Ollama): self.llm = llm self.engine = create_engine(f"mysql+pymysql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}") self.sql_db = SQLDatabase(self.engine) self.table_names = self.sql_db.get_usable_table_names() self.schema_table_mapping = {} self.init_schema_table_mapping() self.modules = self.prepare_modules() self.pipeline = self.build_pipeline() def init_schema_table_mapping(self): self.table_infos = [] table_names = set() for table in self.table_names: table_info = TableInfo(table_name=table, table_summary=TABLE_SUMMARY[table]) self.table_infos.append(table_info) # 摘要表名: 真实表名 self.schema_table_mapping[table_info.table_name] = table def prepare_modules(self): modules = {} # input modules["input"] = InputComponent() # table retriever table_obj_index_path = "/home/purui/projects/chatbot/kb/sql/table_obj_index" retriever = create_table_obj_retriever( index_path=table_obj_index_path, table_infos=self.table_infos, sql_db=self.sql_db, schema_table_mapping=self.schema_table_mapping ) modules["table_retriever"] = TableRetrieveComponent( retriever=retriever, sql_database=self.sql_db ) # text2sql_prompt text2sql_prompt = DEFAULT_TEXT_TO_SQL_PROMPT.partial_format( dialect=self.engine.dialect.name ) modules["text2sql_prompt"] = text2sql_prompt # text2sql_llm modules["text2sql_llm"] = self.llm # sql output parser modules["sql_output_parser"] = FnComponent(fn=parse_response_to_sql) # sql retriever # modules["sql_retriever"] = SQLRetriever(self.sql_db) modules["sql_retriever"] = CustomSQLRetriever(sql_db=self.sql_db) # response synthesise prompt response_synthesis_prompt_str = ( "Given an input question, synthesize a response from the query results.\n" "Query: {query_str}\n" "SQL: {sql_query}\n" "SQL Response: {context_str}\n" "Response: " ) response_synthesis_prompt = PromptTemplate( response_synthesis_prompt_str, ) modules["response_synthesis_prompt"] = response_synthesis_prompt # response synthesise llm modules["response_synthesis_llm"] = self.llm return modules def build_pipeline(self): qp = QueryPipeline( modules=self.modules, verbose=True, ) # add chains & links qp.add_link("input", "table_retriever", dest_key="query") qp.add_link("input", "text2sql_prompt", dest_key="query_str") qp.add_link("table_retriever", "text2sql_prompt", dest_key="schema") qp.add_chain( ["text2sql_prompt", "text2sql_llm", "sql_output_parser"] ) qp.add_link( "sql_output_parser", "response_synthesis_prompt", dest_key="sql_query" ) qp.add_link("input", "sql_retriever", dest_key="query_str") qp.add_link("sql_output_parser", "sql_retriever", dest_key="sql_query") # custom sql_retriever component：定义is_valid字段，如果执行sql检索有正确返回结果，则is_valid为True 作为sql_retriever -> response_synthesis_prompt的链接条件 # 若is_valid为False，则重新回到text2sql_prompt链路中，重新生成sql qp.add_link( "sql_retriever", "response_synthesis_prompt", dest_key="context_str", condition_fn=lambda x: x["is_valid"] ) qp.add_link("sql_retriever", "text2sql_prompt", src_key="query_str", dest_key="query_str", condition_fn=lambda x: not x["is_valid"]) qp.add_link("input", "response_synthesis_prompt", dest_key="query_str") qp.add_link("response_synthesis_prompt", "response_synthesis_llm") return qp def get_vision(self): net = Network(notebook=True, cdn_resources="in_line", directed=True) net.from_nx(self.pipeline.dag) net.write_html("text2sql_dag.html") with open("text2sql_dag.html", "r") as file: html_content = file.read() # Display the HTML content display(HTML(html_content)) def run(self, query: str): response = self.pipeline.run(query=query) return str(response) class TableInfo(BaseModel): """Information regarding a structured table.""" table_name: str = Field( ..., description="table name (must be underscores and NO spaces)" ) table_summary: str = Field( ..., description="short, concise summary/caption of the table" ) class TableRetrieveComponent(CustomQueryComponent): """Retrieves table information from the database.""" retriever: ObjectRetriever = Field(..., description="Retriever to use for table info") sql_database: SQLDatabase = Field(..., description="SQL engine to use for table info") def _validate_component_inputs( self, input: Dict[str, Any] ) -> Dict[str, Any]: """Validate component inputs during run_component.""" # NOTE: this is OPTIONAL but we show you here how to do validation as an example return input @property def _input_keys(self) -> set: """Input keys dict.""" return {"query"} @property def _output_keys(self) -> set: # can do multi-outputs too return {"output"} def _run_component(self, **kwargs) -> Dict[str, Any]: """Run the component.""" # run logic table_schema = self.retriever.retrieve(kwargs["query"])[0] table_name = table_schema.table_name table_info = TABLE_SUMMARY[table_name] return {"output": table_info} if __name__ == '__main__': sql_pipeline = SQLPipeline(llm=Ollama(model="mannix/llama3.1-8b-abliterated", request_timeout=120)) response = sql_pipeline.run("I want 5 different big tits milf porn with it's title and web url") print(response) # table_retriever = sql_pipeline.modules["table_retriever"] # # result = table_retriever.retrieve("Give me top 5 videos by view count.") # # print(result) # qp = QueryPipeline( # modules={ # "input": InputComponent(), # "table_retriever": TableRetrieveComponent(retriever=table_retriever, sql_database=sql_pipeline.sql_db), # } # ) # qp.add_link("input", "table_retriever", dest_key="query") # response = qp.run(query="Give me top 5 videos by view count.") # print(response)