Update logic.py

logic.py

@@ -3,18 +3,20 @@ import kuzu
 import logging
 import sys
 import os
-from llama_index.graph_stores.kuzu import KuzuGraphStore
-from llama_index.core import (
+
+#import llama_index
+from llama_index.graph_stores import KuzuGraphStore
+from llama_index import (
     SimpleDirectoryReader,
     ServiceContext,
     KnowledgeGraphIndex,
 )
-from llama_index.readers.web import SimpleWebPageReader
-
+from llama_index.readers import SimpleWebPageReader
+from llama_index.indices.loading import load_index_from_storage
 
-from llama_index.llms.openai import OpenAI
+from llama_index.llms import OpenAI
 from IPython.display import Markdown, display
-from llama_index.core.storage.storage_context import StorageContext
+from llama_index.storage.storage_context import StorageContext
 
 from pyvis.network import Network
 import pandas as pd
@@ -22,11 +24,26 @@ import numpy as np
 import plotly.express as px
 import umap
 
-def get_index(links,token):
-    #os.environ["OPENAI_API_KEY"] = userdata.get('oai')
+def load_index(token,name):
+    os.environ["OPENAI_API_KEY"] = token
+    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+    
+    db = kuzu.Database(name+"/kg")
+    graph_store = KuzuGraphStore(db)
+    llm = OpenAI(temperature=0, model="gpt-3.5-turbo",api_key=token)
+
+    service_context = ServiceContext.from_defaults(llm=llm, chunk_size=512)
+    storage_context = StorageContext.from_defaults(graph_store=graph_store,persist_dir=name+"/storage")
+    index = load_index_from_storage(storage_context=storage_context,service_context=service_context)
+    return index
+
+
+def get_index(links,token,name):
+    os.mkdir(name) 
+    os.environ["OPENAI_API_KEY"] = token
     logging.basicConfig(stream=sys.stdout, level=logging.INFO)
     
-    db = kuzu.Database("kg1")
+    db = kuzu.Database(name+"/kg")
     graph_store = KuzuGraphStore(db)
     
     
@@ -36,24 +53,25 @@ def get_index(links,token):
     
     llm = OpenAI(temperature=0, model="gpt-3.5-turbo",api_key=token)
     service_context = ServiceContext.from_defaults(llm=llm, chunk_size=512)
-    
     storage_context = StorageContext.from_defaults(graph_store=graph_store)
     
     # NOTE: can take a while!
     index = KnowledgeGraphIndex.from_documents(documents=documents,
-                                               max_triplets_per_chunk=5,
+                                               max_triplets_per_chunk=2,
                                                storage_context=storage_context,
                                                service_context=service_context,
                                                show_progress=True,
                                                include_embeddings=True)
+    index.storage_context.persist(name+"/storage")
+   
 
     return index
 
 def get_network_graph(index):
     g = index.get_networkx_graph()
-    net = Network(notebook=True, cdn_resources="in_line", directed=True)
+    net = Network(directed=True)
     net.from_nx(g)
-    net.show("kuzugraph_draw3.html")
+    # net.show("kuzugraph_draw3.html")
     net.save_graph("kuzugraph_draw3.html")
 
 
@@ -84,11 +102,15 @@ def get_visualize_embeddings(embedding_series, n_neighbors=15, min_dist=0.1, n_c
     return fig
 
 
+def query_model(index,user_query):
+    query_engine = index.as_query_engine(
+    include_text=True,
+    response_mode="tree_summarize",
+    embedding_mode="hybrid",
+    similarity_top_k=5,
+)
 
-
-
-
-
-
+    response = query_engine.query(user_query)
+    return response