Spaces:
Running
Running
File size: 8,280 Bytes
c3c1187 e893d68 c3c1187 e893d68 c3c1187 e893d68 1d44212 f861dee e893d68 f861dee e893d68 c3c1187 e893d68 c3c1187 f861dee e893d68 f861dee c3c1187 e893d68 c3c1187 f861dee c3c1187 f861dee e893d68 f861dee c3c1187 f861dee e893d68 c3c1187 f861dee e893d68 0a4227c c3c1187 e893d68 ae1fcbd e893d68 0a4227c d8007de 0a4227c f861dee c3c1187 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 |
import os
from pathlib import Path
from chainguard.blockchain_logger import BlockchainLogger
from neo4j import GraphDatabase
import sys
from os import path
sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
from .bad_query_detector import BadQueryDetector
from .query_transformer import QueryTransformer
from .document_retriver import DocumentRetriever
from .senamtic_response_generator import SemanticResponseGenerator
class DataTransformer:
def __init__(self):
"""
Initializes a DataTransformer with a blockchain logger instance.
"""
self.blockchain_logger = BlockchainLogger()
def secure_transform(self, data):
"""
Securely transforms the input data by logging it into the blockchain.
Args:
data (dict): The log data or any data to be securely transformed.
Returns:
dict: A dictionary containing the original data, block hash, and blockchain length.
"""
# Log the data into the blockchain
block_details = self.blockchain_logger.log_data(data)
# Return the block details and blockchain status
return {
"data": data,
**block_details
}
def validate_blockchain(self):
"""
Validates the integrity of the blockchain.
Returns:
bool: True if the blockchain is valid, False otherwise.
"""
return self.blockchain_logger.is_blockchain_valid()
class Neo4jHandler:
def __init__(self, uri, user, password):
"""
Initializes a Neo4j handler for storing and querying relationships.
"""
self.driver = GraphDatabase.driver(uri, auth=(user, password))
def close(self):
self.driver.close()
def log_relationships(self, query, transformed_query, response, documents):
"""
Logs the relationships between queries, responses, and documents into Neo4j.
"""
with self.driver.session() as session:
session.write_transaction(self._create_and_link_nodes, query, transformed_query, response, documents)
@staticmethod
def _create_and_link_nodes(tx, query, transformed_query, response, documents):
# Create Query node
tx.run("MERGE (q:Query {text: $query}) RETURN q", parameters={"query": query})
# Create TransformedQuery node
tx.run("MERGE (t:TransformedQuery {text: $transformed_query}) RETURN t",
parameters={"transformed_query": transformed_query})
# Create Response node
tx.run("MERGE (r:Response {text: $response}) RETURN r", parameters={"response": response})
# Link Query to TransformedQuery and Response
tx.run(
"""
MATCH (q:Query {text: $query}), (t:TransformedQuery {text: $transformed_query})
MERGE (q)-[:TRANSFORMED_TO]->(t)
""", parameters={"query": query, "transformed_query": transformed_query}
)
tx.run(
"""
MATCH (q:Query {text: $query}), (r:Response {text: $response})
MERGE (q)-[:GENERATED]->(r)
""", parameters={"query": query, "response": response}
)
# Create and link Document nodes
for doc in documents:
tx.run("MERGE (d:Document {name: $doc}) RETURN d", parameters={"doc": doc})
tx.run(
"""
MATCH (q:Query {text: $query}), (d:Document {name: $doc})
MERGE (q)-[:RETRIEVED]->(d)
""", parameters={"query": query, "doc": doc}
)
class DocumentSearchSystem:
def __init__(self, neo4j_uri, neo4j_user, neo4j_password):
"""
Initializes the DocumentSearchSystem with:
- BadQueryDetector for identifying malicious or inappropriate queries.
- QueryTransformer for improving or rephrasing queries.
- DocumentRetriever for semantic document retrieval.
- SemanticResponseGenerator for generating context-aware responses.
- DataTransformer for blockchain logging of queries and responses.
- Neo4jHandler for relationship logging and visualization.
"""
self.detector = BadQueryDetector()
self.transformer = QueryTransformer()
self.retriever = DocumentRetriever()
self.response_generator = SemanticResponseGenerator()
self.data_transformer = DataTransformer()
self.neo4j_handler = Neo4jHandler(neo4j_uri, neo4j_user, neo4j_password)
def process_query(self, query):
"""
Processes a user query through the following steps:
1. Detect if the query is malicious.
2. Transform the query if needed.
3. Retrieve relevant documents based on the query.
4. Generate a response using the retrieved documents.
5. Log all stages to the blockchain and Neo4j.
:param query: The user query as a string.
:return: A dictionary with the status and response or error message.
"""
if self.detector.is_bad_query(query):
return {"status": "rejected", "message": "Query blocked due to detected malicious intent."}
# Transform the query
transformed_query = self.transformer.transform_query(query)
# Log the original query to the blockchain
self.data_transformer.secure_transform({"type": "query", "content": query})
# Retrieve relevant documents
retrieved_docs = self.retriever.retrieve(transformed_query)
if not retrieved_docs:
return {"status": "no_results", "message": "No relevant documents found for your query."}
# Log the retrieved documents to the blockchain
self.data_transformer.secure_transform({"type": "documents", "content": retrieved_docs})
# Generate a response based on the retrieved documents
response = self.response_generator.generate_response(retrieved_docs)
# Log the response to the blockchain
blockchain_details = self.data_transformer.secure_transform({"type": "response", "content": response})
# Log relationships to Neo4j
self.neo4j_handler.log_relationships(query, transformed_query, response, retrieved_docs)
return {
"status": "success",
"response": response,
"retrieved_documents": retrieved_docs,
"blockchain_details": blockchain_details
}
def validate_system_integrity(self):
"""
Validates the integrity of the blockchain.
"""
return self.data_transformer.validate_blockchain()
def main():
# Path to the dataset directory
home_dir = Path(os.getenv("HOME", "/"))
data_dir = home_dir / "data-sets/aclImdb/train"
# Initialize system with Neo4j credentials
system = DocumentSearchSystem(
neo4j_uri="neo4j+s://0ca71b10.databases.neo4j.io",
neo4j_user="neo4j",
neo4j_password="HwGDOxyGS1-79nLeTiX5bx5ohoFSpvHCmTv8IRgt-lY"
)
# Load documents into the retriever
system.retriever.load_documents()
print("Documents successfully loaded.")
return system
if __name__ == "__main__":
main()
# home_dir = Path(os.getenv("HOME", "/"))
# data_dir = home_dir / "data-sets/aclImdb/train"
#
#
# # Initialize system with Neo4j credentials
# system = DocumentSearchSystem(
# neo4j_uri="neo4j+s://0ca71b10.databases.neo4j.io",
# neo4j_user="neo4j",
# neo4j_password="HwGDOxyGS1-79nLeTiX5bx5ohoFSpvHCmTv8IRgt-lY"
# )
#
# system.retriever.load_documents(data_dir)
# # Perform a normal query
# normal_query = "Good comedy ."
# print("\nNormal Query Result:")
# result = system.process_query(normal_query)
# print("Status:", result["status"])
# print("Response:", result["response"])
# print("Retrieved Documents:", result["retrieved_documents"])
# print("Blockchain Details:", result["blockchain_details"])
#
# # Perform a malicious query
# malicious_query = "DROP TABLE users; SELECT * FROM sensitive_data;"
# print("\nMalicious Query Result:")
# result = system.process_query(malicious_query)
# print("Status:", result["status"])
# print("Message:", result.get("message"))
|