Update app.py

app.py

@@ -4,18 +4,16 @@ from chainlit.input_widget import Switch
 import google.generativeai as gemini_client
 from qdrant_client import QdrantClient
 
-import vanna
-from vanna.vannadb import VannaDB_VectorStore
-from vanna.google import GoogleGeminiChat
-
 import textwrap
 import os
+import requests
+import json
 
 #API KEY
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 QDRANT_URL = os.environ.get("QDRANT_URL")
 QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY")
-VANNA_MODEL = os.environ.get("VANNA_MODEL")
+VANNA_API_URL = os.environ.get("VANNA_API_URL")
 VANNA_API_KEY = os.environ.get("VANNA_API_KEY")
 
 #CONNECTION
@@ -25,22 +23,45 @@ client = QdrantClient(
 )
 gemini_client.configure(api_key=GEMINI_API_KEY)
 
-class MyVanna(VannaDB_VectorStore, GoogleGeminiChat):
-    def __init__(self, config=None):
-        MY_VANNA_MODEL = VANNA_MODEL
-        VannaDB_VectorStore.__init__(self, vanna_model=VANNA_MODEL, vanna_api_key=VANNA_API_KEY, config=config)
-        GoogleGeminiChat.__init__(self, config={'api_key': GEMINI_API_KEY, 'model': 'gemini-1.5-flash',"language":"id-ID"})
+def call_vanna_api(message):
+    headers = {
+        "Content-Type": "application/json",
+        "VANNA-API-KEY": VANNA_API_KEY
+    }
+    data = {
+        "message": message,
+        "user_email": "[email protected]",  # Ganti dengan email pengguna
+        "agent_id": "landw",
+        "acceptable_responses": ["text"],
+    }
 
-# Error handling for session disconnection
-async def handle_request(self, scope, receive, send):
     try:
-        socket = self._get_socket(sid)
-    except KeyError:
-        print("Session is disconnected")
-        return
-    # Continue with the rest of the request handling
+        response = requests.post(VANNA_API_URL, headers=headers, data=json.dumps(data), stream=True)
+        response.raise_for_status()
+
+        results = []
+        for line in response.iter_lines():
+            if line:
+                decoded_line = line.decode('utf-8')
+                if decoded_line.startswith("data:"):
+                    data_string = decoded_line[5:].strip()
+                    try:
+                        data = json.loads(data_string)
+                        if data['type'] == 'text':
+                            results.append(data['text'])  # Append the text response to the results list
+                        elif data['type'] == 'end':
+                            # Return only the last 'text' type response
+                            return results[-1] if results else "No valid text response found"
+                        elif data['type'] == 'error':
+                            results.append(f"**Error:**\n{data['error']}")
+                    except json.JSONDecodeError as e:
+                        results.append(f"Error decoding JSON: {e} - Original data: {data_string}")
+        return "\n\n".join(results)
+    except requests.exceptions.RequestException as e:
+        return f"An error occurred: {e}"
+    except Exception as e:
+        return f"An unexpected error occurred: {e}"
 
-vn = MyVanna()
 
 @cl.set_chat_profiles
 async def chat_profile():
@@ -106,52 +127,10 @@ async def main(message: cl.Message):
     if cl.user_session.get("chat_profile") == "Databot LAN":
         human_query = message.content
 
-        vn.connect_to_postgres(host=os.environ.get("DB_HOST"), dbname=os.environ.get("DB_NAME"), user=os.environ.get("DB_USER"), password=os.environ.get("DB_PASS"), port=os.environ.get("DB_PORT"))
-
-        def gen_query(human_query: str):
-            sql_query = vn.generate_sql(human_query)
-            return sql_query
-
-        def execute_query(query):
-            df = vn.run_sql(query)
-
-            return df
-
-        def human_answer(human_query,df):
-            answer = vn.generate_summary(human_query, df)
-
-            return answer
-
-        def plot(human_query, sql, df):
-            plotly_code = vn.generate_plotly_code(question=human_query, sql=sql, df=df)
-            fig = vn.get_plotly_figure(plotly_code=plotly_code, df=df)
+        # Panggil API Vanna AI untuk mendapatkan respons
+        vanna_response = call_vanna_api(human_query)
 
-            elements = [cl.Plotly(name="chart", figure=fig, display="inline", root=False)]
-            return elements
-            
-
-        settings = await cl.ChatSettings(
-            [
-                Switch(
-                    id="is_chart",
-                    label="Tampilkan grafik",
-                    initial=False
-                    )
-            ]
-        ).send()
-
-        sql_query = gen_query(human_query)
-        df = execute_query(sql_query)
-
-        answer_human = human_answer(human_query, df)
-
-        if settings["is_chart"]:
-            elements = plot(human_query, sql_query, df)
-            await cl.Message(
-                content=human_query, elements=elements, author="Databot LANRI").send()
-
-        await cl.Message(
-            content=f"{answer_human}").send()
+        await cl.Message(content=vanna_response).send()
 
     elif cl.user_session.get("chat_profile") == "Infobot LAN":
         ask = message.content
@@ -187,7 +166,7 @@ async def main(message: cl.Message):
         prompt = make_prompt(question=ask,
                          relevant_passage = search_result[0].payload['answer'],
                          data_source = search_result[0].payload['pasal'] + '  ' + search_result[0].payload['peraturan'])
-        model = gemini_client.GenerativeModel('models/gemini-1.5-flash')
+        model = gemini_client.GenerativeModel('models/gemini-2.0-flash')
         answer = model.generate_content(prompt)
 
         await cl.Message(