| from haystack import Pipeline
|
| from haystack.components.builders import PromptBuilder
|
| from haystack.components.generators.openai import OpenAIGenerator
|
| from haystack.components.routers import ConditionalRouter
|
|
|
| from functions import SQLiteQuery
|
|
|
| from typing import List
|
| import sqlite3
|
|
|
| import os
|
| from getpass import getpass
|
| from dotenv import load_dotenv
|
|
|
| load_dotenv()
|
|
|
| if "OPENAI_API_KEY" not in os.environ:
|
| os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key:")
|
|
|
| from haystack.components.builders import PromptBuilder
|
| from haystack.components.generators import OpenAIGenerator
|
|
|
| llm = OpenAIGenerator(model="gpt-4o")
|
| def rag_pipeline_func(queries: str, columns: str, session_hash):
|
| sql_query = SQLiteQuery(f'data_source_{session_hash}.db')
|
|
|
| connection = sqlite3.connect(f'data_source_{session_hash}.db')
|
| cur=connection.execute('select * from data_source')
|
| columns = [i[0] for i in cur.description]
|
| cur.close()
|
|
|
|
|
| prompt = PromptBuilder(template="""Please generate an SQL query. The query should answer the following Question: {{question}};
|
| If the question cannot be answered given the provided table and columns, return 'no_answer'
|
| The query is to be answered for the table is called 'data_source' with the following
|
| Columns: {{columns}};
|
| Answer:""")
|
|
|
| routes = [
|
| {
|
| "condition": "{{'no_answer' not in replies[0]}}",
|
| "output": "{{replies}}",
|
| "output_name": "sql",
|
| "output_type": List[str],
|
| },
|
| {
|
| "condition": "{{'no_answer' in replies[0]}}",
|
| "output": "{{question}}",
|
| "output_name": "go_to_fallback",
|
| "output_type": str,
|
| },
|
| ]
|
|
|
| router = ConditionalRouter(routes)
|
|
|
| fallback_prompt = PromptBuilder(template="""User entered a query that cannot be answered with the given table.
|
| The query was: {{question}} and the table had columns: {{columns}}.
|
| Let the user know why the question cannot be answered""")
|
| fallback_llm = OpenAIGenerator(model="gpt-4")
|
|
|
| conditional_sql_pipeline = Pipeline()
|
| conditional_sql_pipeline.add_component("prompt", prompt)
|
| conditional_sql_pipeline.add_component("llm", llm)
|
| conditional_sql_pipeline.add_component("router", router)
|
| conditional_sql_pipeline.add_component("fallback_prompt", fallback_prompt)
|
| conditional_sql_pipeline.add_component("fallback_llm", fallback_llm)
|
| conditional_sql_pipeline.add_component("sql_querier", sql_query)
|
|
|
| conditional_sql_pipeline.connect("prompt", "llm")
|
| conditional_sql_pipeline.connect("llm.replies", "router.replies")
|
| conditional_sql_pipeline.connect("router.sql", "sql_querier.queries")
|
| conditional_sql_pipeline.connect("router.go_to_fallback", "fallback_prompt.question")
|
| conditional_sql_pipeline.connect("fallback_prompt", "fallback_llm")
|
|
|
| print("RAG PIPELINE FUNCTION")
|
| result = conditional_sql_pipeline.run({"prompt": {"question": queries,
|
| "columns": columns},
|
| "router": {"question": queries},
|
| "fallback_prompt": {"columns": columns}})
|
|
|
| if 'sql_querier' in result:
|
| reply = result['sql_querier']['results'][0]
|
| elif 'fallback_llm' in result:
|
| reply = result['fallback_llm']['replies'][0]
|
| else:
|
| reply = result["llm"]["replies"][0]
|
|
|
| print("reply content")
|
| print(reply.content)
|
|
|
| return {"reply": reply.content} |
