Update app.py

app.py

@@ -3,7 +3,6 @@ import openai
 import json
 from graphviz import Digraph
 import base64
-from io import BytesIO
 from PIL import Image
 
 def generate_knowledge_graph(api_key, user_input):
@@ -11,24 +10,38 @@ def generate_knowledge_graph(api_key, user_input):
     openai.api_key = api_key
 
     print("Making API call to OpenAI...")
-    completion = openai.Completion.create(
-        engine="text-davinci-002",
-        prompt=f"Help me understand the following by describing it as a detailed knowledge graph: {user_input}",
-        max_tokens=100
+    completion = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo-16k",
+        messages=[
+            {
+                "role": "user",
+                "content": f"Help me understand following by describing as a detailed knowledge graph: {user_input}",
+            }
+        ],
+        functions=[
+            {
+                "name": "knowledge_graph",
+                "description": "Generate a knowledge graph with entities and relationships.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "metadata": {"type": "object"},
+                        "nodes": {"type": "array"},
+                        "edges": {"type": "array"}
+                    },
+                    "required": ["nodes", "edges"]
+                }
+            }
+        ],
+        function_call={"name": "knowledge_graph"}
     )
 
     print("Received response from OpenAI.")
-    response_data = completion.choices[0].text
+    response_data = completion.choices[0]["message"]["function_call"]["arguments"]
     print(f"Response data: {response_data}")
 
-    # For demonstration, let's assume the response_data is a JSON string that can be converted to a dictionary.
-    # You'll need to write code to interpret the text-based response to generate this dictionary.
     print("Converting response to JSON...")
-    try:
-        response_dict = json.loads(response_data)
-    except json.JSONDecodeError:
-        print("Failed to decode JSON. Using empty dictionary as a fallback.")
-        response_dict = {}
+    response_dict = json.loads(response_data)
 
     print("Generating knowledge graph using Graphviz...")
     dot = Digraph(comment="Knowledge Graph")