Test

kg_builder/README.md

@@ -13,6 +13,7 @@ This project builds and queries knowledge graphs from Wikipedia articles using t
 - pip and virtualenv (optional)
 - Access to a Neo4j database
 - OpenAI API key
+- Extra change
 
 ## Installation
 1. **Clone the repository**:

kg_builder/src/graph_creation.py

@@ -1,10 +1,8 @@
 from langchain_community.document_loaders import WikipediaLoader
 from langchain.text_splitter import TokenTextSplitter
 from knowledge_graph_builder import extract_and_store_graph
-from langchain.schema import Document
 from dotenv import load_dotenv
 from tqdm import tqdm
-import os
 
 # Load environment variables
 load_dotenv()

kg_builder/src/knowledge_graph_builder.py

@@ -13,13 +13,6 @@ from models import Node, Relationship, KnowledgeGraph
 from utils import map_to_base_node, map_to_base_relationship
 from api_connections import get_extraction_chain
 
-from langchain.chains.openai_functions import (
-    create_openai_fn_chain,
-    create_structured_output_runnable,
-    create_structured_output_chain,
-)
-
-
 def extract_and_store_graph(
     document: Document,
     category: str,

kg_builder/src/main.py

@@ -1,43 +1,42 @@
 import os
 from openai import OpenAI
 from api_connections import get_graph_connection
-from knowledge_graph_builder import extract_and_store_graph
 from query_graph import query_knowledge_graph
-from langchain_community.document_loaders import WikipediaLoader
-from langchain.text_splitter import TokenTextSplitter
-from tqdm import tqdm
+
+import openai
 
 def get_llm():
     api_key = os.getenv("OPENAI_API_KEY")
     if not api_key:
         raise ValueError("No OpenAI API key found in environment variables.")
-    return OpenAI(api_key=api_key)
+    openai.api_key = api_key
+    return openai
 
 def classify_query(query):
     llm = get_llm()
-    response = llm.Completion.create(
-        model="text-davinci-003",  # Consider updating to the latest model as necessary
-        prompt=f"Classify the following query into 'Chemotherapy' or 'Traffic Law': {query}",
-        max_tokens=60
-    )
-    return response.choices[0].text.strip()
+    try:
+        response = llm.ChatCompletion.create(
+            model="gpt-3.5-turbo-16k",
+            messages=[{"role": "user", "content": f"Classify the following query into 'Chemotherapy' or 'Traffic Law': {query}"}],
+            max_tokens=60
+        )
+        return response['choices'][0]['message']['content'].strip()
+    except Exception as e:
+        print(f"Error during classification: {e}")
+        return None
+
 
 def main():
     print("Starting the script...")
-
     # Get user query
-    query = input("Please enter your query: ")
-    
+    query = input("Please enter your query: ")    
     # Classify the query
     category = classify_query(query)
-    print(f"Query classified into category: {category}")
-    
+    print(f"Query classified into category: {category}")    
     # Get the correct graph connection
-    graph = get_graph_connection(category)
-    
+    graph = get_graph_connection(category)    
     # Query the correct graph
     result = query_knowledge_graph(graph, query)
     print(f"Query result: {result}")
-
 if __name__ == "__main__":
     main()

kg_builder/src/models.py

@@ -1,7 +1,6 @@
 from langchain_community.graphs.graph_document import (
     Node as BaseNode,
     Relationship as BaseRelationship,
-    GraphDocument,
 )
 from langchain.schema import Document
 from typing import List, Dict, Any, Optional

kg_builder/src/query_graph.py

@@ -1,6 +1,5 @@
 from langchain.chains import GraphCypherQAChain
 from langchain_openai import ChatOpenAI
-from api_connections import graph  # Importing 'graph' from 'api_connections.py'
 
 def query_knowledge_graph(graph, query):
     print("Refreshing the graph schema...")

kg_builder/src/utils.py

@@ -1,7 +1,6 @@
 from langchain_community.graphs.graph_document import (
     Node as BaseNode,
     Relationship as BaseRelationship,
-    GraphDocument,
 )
 from models import Node, Relationship