Merge pull request #4 from FalkorDB/verification-token

Verification token

Author: gkorland

.env.template

@@ -0,0 +1,8 @@
+FALKORDB_HOST=<host>
+FALKORDB_PORT=<port>
+FALKORDB_USERNAME=<username>
+FALKORDB_PASSWORD=<password>
+OPENAI_API_KEY=<openai_api_key>
+GOOGLE_API_KEY=<google_api_key>
+SECRET_TOKEN=<secret_token>
+

.gitignore

@@ -0,0 +1 @@
+.env

code_graph/app.py

@@ -2,6 +2,7 @@
 import datetime
 from code_graph import *
 from typing import Optional
+from functools import wraps
 from falkordb import FalkorDB
 from dotenv import load_dotenv
 from urllib.parse import urlparse
@@ -17,10 +18,26 @@
                     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# Function to verify the token
+SECRET_TOKEN = os.getenv('SECRET_TOKEN')
+def verify_token(token):
+    return token == SECRET_TOKEN
+
+# Decorator to protect routes with token authentication
+def token_required(f):
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        token = request.headers.get('Authorization')  # Get token from header
+        if not token or not verify_token(token):
+            return jsonify(message="Unauthorized"), 401
+        return f(*args, **kwargs)
+    return decorated_function
+
 def create_app():
     app = Flask(__name__)
 
     @app.route('/graph_entities', methods=['GET'])
+    @token_required  # Apply token authentication decorator
     def graph_entities():
         """
         Endpoint to fetch sub-graph entities from a given repository.
@@ -64,6 +81,7 @@ def graph_entities():
 
 
     @app.route('/get_neighbors', methods=['GET'])
+    @token_required  # Apply token authentication decorator
     def get_neighbors():
         """
         Endpoint to get neighbors of a specific node in the graph.
@@ -117,6 +135,7 @@ def get_neighbors():
 
 
     @app.route('/process_repo', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def process_repo():
         """
         Process a GitHub repository.
@@ -151,6 +170,7 @@ def process_repo():
         return jsonify(response), 200
 
     @app.route('/process_local_repo', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def process_local_repo():
         # Get JSON data from the request
         data = request.get_json()
@@ -182,6 +202,7 @@ def process_local_repo():
         return jsonify(response), 200
 
     @app.route('/process_code_coverage', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def process_code_coverage():
         """
         Endpoint to process code coverage data for a given repository.
@@ -217,6 +238,7 @@ def process_code_coverage():
 
 
     @app.route('/switch_commit', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def switch_commit():
         """
         Endpoint to switch a repository to a specific commit.
@@ -250,6 +272,7 @@ def switch_commit():
         return jsonify(response), 200
 
     @app.route('/auto_complete', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def auto_complete():
         """
         Endpoint to process auto-completion requests for a repository based on a prefix.
@@ -288,6 +311,7 @@ def auto_complete():
 
 
     @app.route('/list_repos', methods=['GET'])
+    @token_required  # Apply token authentication decorator
     def list_repos():
         """
         Endpoint to list all available repositories.
@@ -309,6 +333,7 @@ def list_repos():
 
 
     @app.route('/list_commits', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def list_commits():
         """
         Endpoint to list all commits of a specified repository.
@@ -346,6 +371,7 @@ def list_commits():
 
 
     @app.route('/repo_info', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def repo_info():
         """
         Endpoint to retrieve information about a specific repository.
@@ -390,6 +416,7 @@ def repo_info():
         return jsonify(response), 200
 
     @app.route('/find_paths', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def find_paths():
         """
         Finds all paths between a source node (src) and a destination node (dest) in the graph.
@@ -445,6 +472,7 @@ def find_paths():
 
 
     @app.route('/unreachable', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def unreachable_entities():
         """
         Endpoint to retrieve unreachable entities in the graph.
@@ -478,6 +506,7 @@ def unreachable_entities():
         return jsonify(response), 200
 
     @app.route('/chat', methods=['POST'])
+    @token_required  # Apply token authentication decorator
     def chat():
         # Get JSON data from the request
         data = request.get_json()

code_graph/llm.py

@@ -1,7 +1,7 @@
 import logging
 
 from graphrag_sdk.models.openai import OpenAiGenerativeModel
-from graphrag_sdk.models.gemini import GeminiGenerativeModel
+#from graphrag_sdk.models.gemini import GeminiGenerativeModel
 
 from graphrag_sdk import (
     Ontology,
@@ -193,14 +193,14 @@ def _create_kg_agent(repo_name: str):
     global ontology
 
     openapi_model    = OpenAiGenerativeModel("gpt-4o")
-    gemini_model     = GeminiGenerativeModel("gemini-1.5-flash-001")
-    gemini_model_pro = GeminiGenerativeModel("gemini-1.5-pro")
+    #gemini_model     = GeminiGenerativeModel("gemini-1.5-flash-001")
+    #gemini_model_pro = GeminiGenerativeModel("gemini-1.5-pro")
 
     #ontology = _define_ontology()
     code_graph_kg = KnowledgeGraph(
         name=repo_name,
         ontology=ontology,
-        model_config=KnowledgeGraphModelConfig.with_model(gemini_model),
+        model_config=KnowledgeGraphModelConfig.with_model(openapi_model),
     )
 
     return code_graph_kg.chat_session()