import chainlit as cl from rag_implementation import MistralRAGSystem ## Initialize RAG system rag_system = MistralRAGSystem() ## Pre-populate knowledge graph with some initial data definitialize_knowledge_base(): knowledge_items = [ { "id": "ai_basics", "content": "Artificial Intelligence is a broad field of computer science focused on creating intelligent machines that can simulate human-like thinking and learning capabilities.", "metadata": {"category": "introduction", "difficulty": "beginner"} }, { "id": "ml_fundamentals", "content": "Machine Learning is a subset of AI that enables systems to learn and improve from experience without being explicitly programmed, using algorithms that can learn from and make predictions or decisions based on data.", "metadata": {"category": "core_concept", "difficulty": "intermediate"} } ] for item in knowledge_items: rag_system.add_knowledge(item["id"], item["content"], item["metadata"]) rag_system ## Initialize knowledge base initialize_knowledge_base() @cl.on_chat_start asyncdefstart(): await cl.Message(content="RAG System with Mistral is ready! How can I help you today?").send() @cl.on_message asyncdefmain(message: cl.Message): # Check if the message is a knowledge addition command if message.content.startswith("/add_knowledge"): # Parse the message to extract node_id and content parts = message.content.split(maxsplit=3) iflen(parts) < 3: await cl.Message(content="Usage: /add_knowledge <node_id> <content>").send() return node_id, content = parts[1], parts[2] rag_system.add_knowledge(node_id, content) await cl.Message(content=f"Added knowledge node: {node_id}").send() return # Regular query processing # Augment the query with relevant context augmented_query = rag_system.augment_query(message.content) # Generate response response = rag_system.generate_response(augmented_query) # Send the response back to the user await cl.Message(content=response).send()
import os from dotenv import load_dotenv import requests from graph_embedding import KnowledgeGraphRAG classMistralRAGSystem: def__init__(self): # Load environment variables load_dotenv() # Get Hugging Face API key from environment variable self.api_key = os.getenv('MISTRAL_API_KEY') ifnotself.api_key: raise ValueError("HUGGINGFACE_API_KEY must be set in .env file") # Default model (corrected name) self.model = "mistralai/Mistral-7B-v0.1" # Initialize Knowledge Graph self.knowledge_graph = KnowledgeGraphRAG()
defaugment_query(self, query: str) -> str: """ Augment the query with relevant context from the knowledge graph Args: query (str): Original user query Returns: str: Augmented query with additional context """ # Retrieve similar nodes similar_nodes = self.knowledge_graph.retrieve_similar_nodes(query) # If similar_nodes is a list, iterate over it directly context = "\n".join([str(doc) for doc in similar_nodes]) # Create a structured prompt with context augmented_prompt = f""" #Context Information: #{context} Based on the provided context and your extensive knowledge, please answer the following query comprehensively: Query: {query} Response: """ return augmented_prompt defgenerate_response(self, augmented_query: str) -> str: """ Generate response using Hugging Face API for Mistral model Args: augmented_query (str): Augmented query with context Returns: str: Generated response """ try: # Prepare headers with the Hugging Face API key headers = { 'Authorization': f'Bearer {self.api_key}', 'Content-Type': 'application/json' } # Prepare payload payload = { 'inputs': augmented_query } # Hugging Face Inference API endpoint for Mistral model url = f'https://api-inference.huggingface.co/models/{self.model}' # Make the POST request to generate a response response = requests.post(url, json=payload, headers=headers) # Check if the request was successful if response.status_code == 200: #return response.json()[0]['generated_text'] generated_text = response.json()[0]['generated_text'] print("Raw response:", response.json()) start_index = generated_text.find("Response:") + len("Response:") response_without_context = generated_text[start_index:].strip() return response_without_context else: returnf"Error: {response.status_code} - {response.text}" except Exception as e: returnf"An error occurred: {str(e)}" defadd_knowledge(self, node_id: str, content: str, metadata: dict = None): """ Add knowledge to the graph Args: node_id (str): Unique node identifier content (str): Node content metadata (dict, optional): Additional metadata """ self.knowledge_graph.add_node(node_id, content, metadata)
import networkx as nx import matplotlib.pyplot as plt from sentence_transformers import SentenceTransformer import chromadb from chromadb.config import DEFAULT_TENANT, DEFAULT_DATABASE, Settings from typing importList, Dict, Any classKnowledgeGraphRAG: def__init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"): # Initialize embedding model self.embedding_model = SentenceTransformer(model_name) # Initialize graph self.graph = nx.DiGraph() self.chroma_client = chromadb.PersistentClient( path="test", settings=Settings(), tenant=DEFAULT_TENANT, database=DEFAULT_DATABASE, ) self.collection = self.chroma_client.create_collection(name="knowledge_base3") defadd_node(self, node_id: str, content: str, metadata: Dict[str, Any] = None): """ Add a node to the knowledge graph and embed its content Args: node_id (str): Unique identifier for the node content (str): Text content of the node metadata (dict, optional): Additional metadata for the node """ # Add to networkx graph self.graph.add_node(node_id, content=content, metadata=metadata or {}) # Generate embedding embedding = self.embedding_model.encode(content).tolist() # Ensure metadata is a non-empty dictionary metadata = metadata or {} # Add to ChromaDB self.collection.add( ids=[node_id], embeddings=[embedding], documents=[content], metadatas=[metadata] # Ensure that the metadata is a valid dictionary ) defadd_edge(self, source: str, target: str, relationship: str = None): """ Add a directed edge between two nodes Args: source (str): Source node ID target (str): Target node ID relationship (str, optional): Type of relationship """ self.graph.add_edge(source, target, relationship=relationship) defretrieve_similar_nodes(self, query: str, top_k: int = 3): """ Retrieve most similar nodes to a given query. Args: query (str): Search query top_k (int): Number of top similar nodes to retrieve. Returns: List of most similar nodes. """ # Generate query embedding query_embedding = self.embedding_model.encode(query).tolist() # Get the total number of nodes in the collection total_nodes = self.collection.count() # Adjust top_k if it exceeds the number of available nodes top_k = min(top_k, total_nodes) # Retrieve from ChromaDB results = self.collection.query( query_embeddings=[query_embedding], n_results=top_k ) # Return the documents (already adjusted for n_results) return results.get('documents', []) ## Example usage defcreate_sample_knowledge_graph(): kg = KnowledgeGraphRAG() #persist_directory="./my_knowledge_base_data2" # Add some sample nodes about AI kg.add_node("ai_intro", "人工智能是计算机科学的一个分支") kg.add_node("ml_intro", "机器学习是 AI 的一个子集，专注于从数据中学习") kg.add_node("dl_intro", "深度学习使用具有多个层的神经网络") # Add some relationships kg.add_edge("ai_intro", "ml_intro", "包含") kg.add_edge("ml_intro", "dl_intro", "高级技术") return kg ## For testing if __name__ == "__main__": kg = create_sample_knowledge_graph() kg.visualaze_graph() # Example retrieval results = kg.retrieve_similar_nodes("神经网络") print(results)
粤ICP备14082021号
免费POC，
零成本试错
