Buckets:

tahamajs
/

Sysmem2_in_AI

Files

xet

tahamajs/Sysmem2_in_AI / ComputerAssignments /CA10_knowledge_graphs /integrations /gemini_experimental.py

tahamajs

4 months ago

download

raw

25.2 kB

"""
Gemini Experimental Knowledge Graph System with LangChain and RAG
پروژه تجربی کامل برای Knowledge Graphs با Gemini API، LangChain و RAG
"""

import os
import json
import logging
import asyncio
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict

# LangChain imports
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
from langchain.agents import initialize_agent, AgentType, Tool
from langchain.tools import DuckDuckGoSearchRun
from langchain.vectorstores import Chroma, FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader, PyPDFLoader, WebBaseLoader
from langchain.schema import Document
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from langchain.callbacks import StreamingStdOutCallbackHandler

# Knowledge Graph specific imports
import rdflib
from rdflib import Graph, Namespace, Literal, URIRef
import spacy
from spacy import displacy
import pickle

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class KGEntity:
    """نمایش یک موجودیت در Knowledge Graph"""

    id: str
    name: str
    type: str
    properties: Dict[str, Any]
    embeddings: Optional[List[float]] = None


@dataclass
class KGRelation:
    """نمایش یک رابطه در Knowledge Graph"""

    subject: str
    predicate: str
    object: str
    confidence: float
    metadata: Dict[str, Any]


class GeminiKnowledgeGraphSystem:
    """سیستم کامل Knowledge Graph با Gemini API"""

    def __init__(self, api_key: str, model_name: str = "gemini-1.5-pro"):
        """مقداردهی سیستم"""
        self.api_key = api_key
        self.model_name = model_name

        # Initialize Gemini LLM
        self.llm = ChatGoogleGenerativeAI(
            model=model_name,
            google_api_key=api_key,
            temperature=0.1,
            max_output_tokens=2048,
            convert_system_message_to_human=True,
        )

        # Initialize embeddings
        self.embeddings = GoogleGenerativeAIEmbeddings(
            model="models/embedding-001", google_api_key=api_key
        )

        # Knowledge Graph components
        self.entities = {}
        self.relations = []
        self.graph = nx.DiGraph()
        self.rdf_graph = Graph()
        self.vector_store = None
        self.memory = ConversationBufferWindowMemory(k=5)

        # NLP model for entity extraction
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            logger.warning(
                "spaCy English model not found. Install with: python -m spacy download en_core_web_sm"
            )
            self.nlp = None

        # Search tool
        self.search = DuckDuckGoSearchRun()

        logger.info("Gemini Knowledge Graph System initialized successfully")

    async def extract_entities_and_relations(
        self, text: str
    ) -> Tuple[List[KGEntity], List[KGRelation]]:
        """استخراج موجودیت‌ها و روابط از متن با استفاده از Gemini"""

        extraction_prompt = f"""
        شما یک متخصص استخراج اطلاعات از Knowledge Graph هستید. 
        از متن زیر، تمام موجودیت‌ها و روابط مهم را استخراج کنید.
        متن:
        {text}
        لطفاً نتیجه را در فرمت JSON زیر برگردانید:
        {{
            "entities": [
                {{
                    "id": "unique_id",
                    "name": "نام موجودیت",
                    "type": "نوع موجودیت (PERSON, ORGANIZATION, LOCATION, CONCEPT, etc.)",
                    "properties": {{
                        "description": "توضیحات",
                        "confidence": 0.95
                    }}
                }}
            ],
            "relations": [
                {{
                    "subject": "موجودیت مبدأ",
                    "predicate": "نوع رابطه",
                    "object": "موجودیت مقصد",
                    "confidence": 0.9,
                    "metadata": {{
                        "context": "زمینه رابطه"
                    }}
                }}
            ]
        }}
        فقط JSON خالص برگردانید، بدون توضیحات اضافی.
        """

        try:
            response = await self.llm.ainvoke(extraction_prompt)
            result = json.loads(response.content)

            entities = [KGEntity(**entity) for entity in result.get("entities", [])]
            relations = [
                KGRelation(**relation) for relation in result.get("relations", [])
            ]

            logger.info(
                f"Extracted {len(entities)} entities and {len(relations)} relations"
            )
            return entities, relations

        except Exception as e:
            logger.error(f"Error in entity extraction: {e}")
            return [], []

    async def build_knowledge_graph(self, documents: List[str]) -> Dict[str, Any]:
        """ساخت Knowledge Graph از اسناد"""

        logger.info(f"Building knowledge graph from {len(documents)} documents")

        all_entities = []
        all_relations = []

        for i, doc in enumerate(documents):
            logger.info(f"Processing document {i+1}/{len(documents)}")

            # Extract entities and relations
            entities, relations = await self.extract_entities_and_relations(doc)
            all_entities.extend(entities)
            all_relations.extend(relations)

            # Add to NetworkX graph
            for entity in entities:
                self.graph.add_node(entity.id, **entity.properties)

            for relation in relations:
                self.graph.add_edge(
                    relation.subject,
                    relation.object,
                    predicate=relation.predicate,
                    confidence=relation.confidence,
                    metadata=relation.metadata,
                )

        # Store entities and relations
        self.entities = {entity.id: entity for entity in all_entities}
        self.relations = all_relations

        # Generate embeddings for entities
        await self._generate_entity_embeddings()

        # Create vector store
        await self._create_vector_store(documents)

        stats = {
            "total_entities": len(all_entities),
            "total_relations": len(all_relations),
            "graph_nodes": self.graph.number_of_nodes(),
            "graph_edges": self.graph.number_of_edges(),
        }

        logger.info(f"Knowledge graph built: {stats}")
        return stats

    async def _generate_entity_embeddings(self):
        """تولید embeddings برای موجودیت‌ها"""

        logger.info("Generating entity embeddings")

        for entity_id, entity in self.entities.items():
            try:
                # Create description for embedding
                description = f"{entity.name} is a {entity.type}. {entity.properties.get('description', '')}"

                # Generate embedding
                embedding = await self.embeddings.aembed_query(description)
                entity.embeddings = embedding

            except Exception as e:
                logger.error(f"Error generating embedding for {entity_id}: {e}")

    async def _create_vector_store(self, documents: List[str]):
        """ایجاد vector store از اسناد"""

        logger.info("Creating vector store")

        try:
            # Create documents
            docs = [
                Document(page_content=doc, metadata={"source": f"doc_{i}"})
                for i, doc in enumerate(documents)
            ]

            # Split documents
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000, chunk_overlap=200
            )
            split_docs = text_splitter.split_documents(docs)

            # Create vector store
            self.vector_store = await FAISS.afrom_documents(split_docs, self.embeddings)

            logger.info(f"Vector store created with {len(split_docs)} chunks")

        except Exception as e:
            logger.error(f"Error creating vector store: {e}")

    async def query_knowledge_graph(
        self, query: str, use_rag: bool = True
    ) -> Dict[str, Any]:
        """پرس‌وجو از Knowledge Graph"""

        logger.info(f"Querying knowledge graph: {query}")

        result = {
            "query": query,
            "direct_answers": [],
            "rag_answers": [],
            "graph_paths": [],
            "related_entities": [],
        }

        try:
            # Direct graph query
            if self.graph.number_of_nodes() > 0:
                result["direct_answers"] = await self._query_graph_directly(query)

            # RAG-based query
            if use_rag and self.vector_store:
                result["rag_answers"] = await self._query_with_rag(query)

            # Find graph paths
            result["graph_paths"] = await self._find_relevant_paths(query)

            # Find related entities
            result["related_entities"] = await self._find_related_entities(query)

        except Exception as e:
            logger.error(f"Error querying knowledge graph: {e}")
            result["error"] = str(e)

        return result

    async def _query_graph_directly(self, query: str) -> List[str]:
        """پرس‌وجو مستقیم از گراف"""

        answers = []

        try:
            # Simple keyword matching
            query_lower = query.lower()

            for node in self.graph.nodes():
                node_data = self.graph.nodes[node]
                if any(
                    keyword in str(node_data).lower() for keyword in query_lower.split()
                ):
                    answers.append(f"Found entity: {node} with properties: {node_data}")

            # Find edges containing query terms
            for edge in self.graph.edges(data=True):
                edge_data = edge[2]
                if any(
                    keyword in str(edge_data).lower() for keyword in query_lower.split()
                ):
                    answers.append(
                        f"Found relation: {edge[0]} -> {edge[1]} ({edge_data.get('predicate', 'unknown')})"
                    )

        except Exception as e:
            logger.error(f"Error in direct graph query: {e}")

        return answers

    async def _query_with_rag(self, query: str) -> List[str]:
        """پرس‌وجو با استفاده از RAG"""

        answers = []

        try:
            # Retrieve relevant documents
            docs = await self.vector_store.asimilarity_search(query, k=3)

            # Create RAG chain
            rag_prompt = PromptTemplate(
                input_variables=["context", "question"],
                template="""
                بر اساس اطلاعات زیر به سوال پاسخ دهید:
                
                اطلاعات مرتبط:
                {context}
                
                سوال: {question}
                
                پاسخ:
                """,
            )

            # Combine retrieved documents
            context = "\n\n".join([doc.page_content for doc in docs])

            # Generate answer with Gemini
            prompt = rag_prompt.format(context=context, question=query)
            response = await self.llm.ainvoke(prompt)

            answers.append(response.content)

        except Exception as e:
            logger.error(f"Error in RAG query: {e}")

        return answers

    async def _find_relevant_paths(self, query: str) -> List[Dict[str, Any]]:
        """یافتن مسیرهای مرتبط در گراف"""

        paths = []

        try:
            query_terms = query.lower().split()

            # Find entities matching query terms
            relevant_nodes = []
            for node in self.graph.nodes():
                node_data = self.graph.nodes[node]
                if any(term in str(node_data).lower() for term in query_terms):
                    relevant_nodes.append(node)

            # Find paths between relevant nodes
            if len(relevant_nodes) >= 2:
                for i in range(len(relevant_nodes)):
                    for j in range(i + 1, len(relevant_nodes)):
                        try:
                            path = nx.shortest_path(
                                self.graph, relevant_nodes[i], relevant_nodes[j]
                            )
                            if len(path) > 1:
                                paths.append(
                                    {
                                        "path": path,
                                        "length": len(path) - 1,
                                        "source": relevant_nodes[i],
                                        "target": relevant_nodes[j],
                                    }
                                )
                        except nx.NetworkXNoPath:
                            continue

        except Exception as e:
            logger.error(f"Error finding paths: {e}")

        return paths

    async def _find_related_entities(self, query: str) -> List[Dict[str, Any]]:
        """یافتن موجودیت‌های مرتبط"""

        related = []

        try:
            # Use embeddings for similarity search
            query_embedding = await self.embeddings.aembed_query(query)

            for entity_id, entity in self.entities.items():
                if entity.embeddings:
                    # Calculate cosine similarity
                    similarity = np.dot(query_embedding, entity.embeddings) / (
                        np.linalg.norm(query_embedding)
                        * np.linalg.norm(entity.embeddings)
                    )

                    if similarity > 0.7:  # Threshold
                        related.append(
                            {
                                "entity_id": entity_id,
                                "name": entity.name,
                                "type": entity.type,
                                "similarity": float(similarity),
                                "properties": entity.properties,
                            }
                        )

            # Sort by similarity
            related.sort(key=lambda x: x["similarity"], reverse=True)

        except Exception as e:
            logger.error(f"Error finding related entities: {e}")

        return related[:5]  # Return top 5

    async def create_rag_agent(self) -> Any:
        """ایجاد RAG Agent با Gemini"""

        try:
            # Define tools
            tools = [
                Tool(
                    name="knowledge_graph_search",
                    description="جستجو در Knowledge Graph برای یافتن اطلاعات",
                    func=lambda query: str(self.query_knowledge_graph(query)),
                ),
                Tool(
                    name="web_search",
                    description="جستجو در اینترنت برای اطلاعات جدید",
                    func=self.search.run,
                ),
            ]

            # Create agent
            agent = initialize_agent(
                tools=tools,
                llm=self.llm,
                agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                memory=self.memory,
                verbose=True,
                handle_parsing_errors=True,
            )

            logger.info("RAG Agent created successfully")
            return agent

        except Exception as e:
            logger.error(f"Error creating RAG agent: {e}")
            return None

    def visualize_knowledge_graph(
        self, max_nodes: int = 50, save_path: str = "kg_visualization.png"
    ):
        """تجسم Knowledge Graph"""

        try:
            plt.figure(figsize=(20, 16))

            # Create subgraph if too many nodes
            if self.graph.number_of_nodes() > max_nodes:
                # Get most connected nodes
                centrality = nx.degree_centrality(self.graph)
                top_nodes = sorted(
                    centrality.items(), key=lambda x: x[1], reverse=True
                )[:max_nodes]
                subgraph = self.graph.subgraph([node for node, _ in top_nodes])
            else:
                subgraph = self.graph

            # Layout
            pos = nx.spring_layout(subgraph, k=3, iterations=50)

            # Draw nodes
            nx.draw_networkx_nodes(
                subgraph, pos, node_color="lightblue", node_size=1000, alpha=0.7
            )

            # Draw edges
            nx.draw_networkx_edges(
                subgraph, pos, edge_color="gray", arrows=True, arrowsize=20, alpha=0.6
            )

            # Draw labels
            nx.draw_networkx_labels(subgraph, pos, font_size=8, font_weight="bold")

            plt.title("Knowledge Graph Visualization", fontsize=16, fontweight="bold")
            plt.axis("off")
            plt.tight_layout()
            plt.savefig(save_path, dpi=300, bbox_inches="tight")
            plt.show()

            logger.info(f"Knowledge graph visualization saved to {save_path}")

        except Exception as e:
            logger.error(f"Error visualizing knowledge graph: {e}")

    def save_knowledge_graph(self, filepath: str):
        """ذخیره Knowledge Graph"""

        try:
            kg_data = {
                "entities": {
                    k: {
                        "id": v.id,
                        "name": v.name,
                        "type": v.type,
                        "properties": v.properties,
                        "embeddings": v.embeddings,
                    }
                    for k, v in self.entities.items()
                },
                "relations": [
                    {
                        "subject": r.subject,
                        "predicate": r.predicate,
                        "object": r.object,
                        "confidence": r.confidence,
                        "metadata": r.metadata,
                    }
                    for r in self.relations
                ],
                "graph_data": {
                    "nodes": list(self.graph.nodes(data=True)),
                    "edges": list(self.graph.edges(data=True)),
                },
                "metadata": {
                    "created_at": datetime.now().isoformat(),
                    "total_entities": len(self.entities),
                    "total_relations": len(self.relations),
                },
            }

            with open(filepath, "w", encoding="utf-8") as f:
                json.dump(kg_data, f, ensure_ascii=False, indent=2)

            logger.info(f"Knowledge graph saved to {filepath}")

        except Exception as e:
            logger.error(f"Error saving knowledge graph: {e}")

    def load_knowledge_graph(self, filepath: str):
        """بارگذاری Knowledge Graph"""

        try:
            with open(filepath, "r", encoding="utf-8") as f:
                kg_data = json.load(f)

            # Load entities
            self.entities = {}
            for k, v in kg_data["entities"].items():
                entity = KGEntity(
                    id=v["id"],
                    name=v["name"],
                    type=v["type"],
                    properties=v["properties"],
                    embeddings=v.get("embeddings"),
                )
                self.entities[k] = entity

            # Load relations
            self.relations = []
            for r in kg_data["relations"]:
                relation = KGRelation(
                    subject=r["subject"],
                    predicate=r["predicate"],
                    object=r["object"],
                    confidence=r["confidence"],
                    metadata=r["metadata"],
                )
                self.relations.append(relation)

            # Rebuild NetworkX graph
            self.graph.clear()
            for node, data in kg_data["graph_data"]["nodes"]:
                self.graph.add_node(node, **data)

            for edge in kg_data["graph_data"]["edges"]:
                self.graph.add_edge(edge[0], edge[1], **edge[2])

            logger.info(f"Knowledge graph loaded from {filepath}")

        except Exception as e:
            logger.error(f"Error loading knowledge graph: {e}")


# Example usage and testing functions
async def run_experimental_examples():
    """اجرای مثال‌های تجربی"""

    # Initialize system
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        logger.error("GEMINI_API_KEY not found in environment variables")
        return

    kg_system = GeminiKnowledgeGraphSystem(api_key)

    # Sample documents
    sample_documents = [
        """
        آلبرت انیشتین یک فیزیکدان نظری آلمانی بود که نظریه نسبیت را توسعه داد. 
        او در سال 1879 در اولم، آلمان متولد شد و در سال 1955 در پرینستون، نیوجرسی درگذشت.
        انیشتین جایزه نوبل فیزیک را در سال 1921 دریافت کرد.
        """,
        """
        نظریه نسبیت انیشتین شامل دو بخش اصلی است: نسبیت خاص و نسبیت عام.
        نسبیت خاص در سال 1905 منتشر شد و شامل رابطه E=mc² است.
        نسبیت عام در سال 1915 منتشر شد و گرانش را به عنوان انحنای فضا-زمان توصیف می‌کند.
        """,
        """
        دانشگاه پرینستون یکی از قدیمی‌ترین و معتبرترین دانشگاه‌های آمریکا است.
        انیشتین از سال 1933 تا زمان مرگش در سال 1955 در موسسه مطالعات پیشرفته پرینستون کار کرد.
        این دانشگاه در نیوجرسی واقع شده است.
        """,
    ]

    try:
        # Build knowledge graph
        print("🔨 Building Knowledge Graph...")
        stats = await kg_system.build_knowledge_graph(sample_documents)
        print(f"✅ Knowledge Graph built: {stats}")

        # Query examples
        queries = [
            "انیشتین کی بود؟",
            "نظریه نسبیت چیست؟",
            "پرینستون کجاست؟",
            "ارتباط انیشتین با پرینستون چیست؟",
        ]

        for query in queries:
            print(f"\n🔍 Query: {query}")
            result = await kg_system.query_knowledge_graph(query)

            print("📝 Direct Answers:")
            for answer in result.get("direct_answers", []):
                print(f"  - {answer}")

            print("🤖 RAG Answers:")
            for answer in result.get("rag_answers", []):
                print(f"  - {answer}")

            print("🛤️ Graph Paths:")
            for path in result.get("graph_paths", []):
                print(f"  - {' -> '.join(path['path'])}")

            print("🔗 Related Entities:")
            for entity in result.get("related_entities", []):
                print(
                    f"  - {entity['name']} ({entity['type']}) - Similarity: {entity['similarity']:.3f}"
                )

        # Create RAG Agent
        print("\n🤖 Creating RAG Agent...")
        agent = await kg_system.create_rag_agent()

        if agent:
            print("✅ RAG Agent created successfully")

            # Test agent
            agent_queries = [
                "درباره انیشتین چه می‌دانید؟",
                "آیا اطلاعات جدیدی درباره نظریه نسبیت وجود دارد؟",
            ]

            for query in agent_queries:
                print(f"\n🤖 Agent Query: {query}")
                try:
                    response = await agent.arun(query)
                    print(f"🤖 Agent Response: {response}")
                except Exception as e:
                    print(f"❌ Agent Error: {e}")

        # Visualize knowledge graph
        print("\n📊 Visualizing Knowledge Graph...")
        kg_system.visualize_knowledge_graph()

        # Save knowledge graph
        print("\n💾 Saving Knowledge Graph...")
        kg_system.save_knowledge_graph("experimental_kg.json")

        print("\n🎉 Experimental examples completed successfully!")

    except Exception as e:
        logger.error(f"Error in experimental examples: {e}")


if __name__ == "__main__":
    # Run experimental examples
    asyncio.run(run_experimental_examples())

Xet Storage Details

Size:: 25.2 kB
Xet hash:: 63442c916bb864e4706150de98312022ee4cf0d204c27f9370d943a05d696527

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""
	Gemini Experimental Knowledge Graph System with LangChain and RAG
	پروژه تجربی کامل برای Knowledge Graphs با Gemini API، LangChain و RAG
	"""

	import os
	import json
	import logging
	import asyncio
	from datetime import datetime
	from typing import List, Dict, Any, Optional, Tuple
	from dataclasses import dataclass
	import networkx as nx
	import matplotlib.pyplot as plt
	import numpy as np
	from collections import defaultdict

	# LangChain imports
	from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
	from langchain.chains import RetrievalQA, ConversationalRetrievalChain
	from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
	from langchain.agents import initialize_agent, AgentType, Tool
	from langchain.tools import DuckDuckGoSearchRun
	from langchain.vectorstores import Chroma, FAISS
	from langchain.text_splitter import RecursiveCharacterTextSplitter
	from langchain.document_loaders import TextLoader, PyPDFLoader, WebBaseLoader
	from langchain.schema import Document
	from langchain.prompts import PromptTemplate, ChatPromptTemplate
	from langchain.callbacks import StreamingStdOutCallbackHandler

	# Knowledge Graph specific imports
	import rdflib
	from rdflib import Graph, Namespace, Literal, URIRef
	import spacy
	from spacy import displacy
	import pickle

	# Setup logging
	logging.basicConfig(level=logging.INFO)
	logger = logging.getLogger(__name__)


	@dataclass
	class KGEntity:
	"""نمایش یک موجودیت در Knowledge Graph"""

	id: str
	name: str
	type: str
	properties: Dict[str, Any]
	embeddings: Optional[List[float]] = None


	@dataclass
	class KGRelation:
	"""نمایش یک رابطه در Knowledge Graph"""

	subject: str
	predicate: str
	object: str
	confidence: float
	metadata: Dict[str, Any]


	class GeminiKnowledgeGraphSystem:
	"""سیستم کامل Knowledge Graph با Gemini API"""

	def __init__(self, api_key: str, model_name: str = "gemini-1.5-pro"):
	"""مقداردهی سیستم"""
	self.api_key = api_key
	self.model_name = model_name

	# Initialize Gemini LLM
	self.llm = ChatGoogleGenerativeAI(
	model=model_name,
	google_api_key=api_key,
	temperature=0.1,
	max_output_tokens=2048,
	convert_system_message_to_human=True,
	)

	# Initialize embeddings
	self.embeddings = GoogleGenerativeAIEmbeddings(
	model="models/embedding-001", google_api_key=api_key
	)

	# Knowledge Graph components
	self.entities = {}
	self.relations = []
	self.graph = nx.DiGraph()
	self.rdf_graph = Graph()
	self.vector_store = None
	self.memory = ConversationBufferWindowMemory(k=5)

	# NLP model for entity extraction
	try:
	self.nlp = spacy.load("en_core_web_sm")
	except OSError:
	logger.warning(
	"spaCy English model not found. Install with: python -m spacy download en_core_web_sm"
	)
	self.nlp = None

	# Search tool
	self.search = DuckDuckGoSearchRun()

	logger.info("Gemini Knowledge Graph System initialized successfully")

	async def extract_entities_and_relations(
	self, text: str
	) -> Tuple[List[KGEntity], List[KGRelation]]:
	"""استخراج موجودیت‌ها و روابط از متن با استفاده از Gemini"""

	extraction_prompt = f"""
	شما یک متخصص استخراج اطلاعات از Knowledge Graph هستید.
	از متن زیر، تمام موجودیت‌ها و روابط مهم را استخراج کنید.

	متن:
	{text}

	لطفاً نتیجه را در فرمت JSON زیر برگردانید:
	{{
	"entities": [
	{{
	"id": "unique_id",
	"name": "نام موجودیت",
	"type": "نوع موجودیت (PERSON, ORGANIZATION, LOCATION, CONCEPT, etc.)",
	"properties": {{
	"description": "توضیحات",
	"confidence": 0.95
	}}
	}}
	],
	"relations": [
	{{
	"subject": "موجودیت مبدأ",
	"predicate": "نوع رابطه",
	"object": "موجودیت مقصد",
	"confidence": 0.9,
	"metadata": {{
	"context": "زمینه رابطه"
	}}
	}}
	]
	}}

	فقط JSON خالص برگردانید، بدون توضیحات اضافی.
	"""

	try:
	response = await self.llm.ainvoke(extraction_prompt)
	result = json.loads(response.content)

	entities = [KGEntity(**entity) for entity in result.get("entities", [])]
	relations = [
	KGRelation(**relation) for relation in result.get("relations", [])
	]

	logger.info(
	f"Extracted {len(entities)} entities and {len(relations)} relations"
	)
	return entities, relations

	except Exception as e:
	logger.error(f"Error in entity extraction: {e}")
	return [], []

	async def build_knowledge_graph(self, documents: List[str]) -> Dict[str, Any]:
	"""ساخت Knowledge Graph از اسناد"""

	logger.info(f"Building knowledge graph from {len(documents)} documents")

	all_entities = []
	all_relations = []

	for i, doc in enumerate(documents):
	logger.info(f"Processing document {i+1}/{len(documents)}")

	# Extract entities and relations
	entities, relations = await self.extract_entities_and_relations(doc)
	all_entities.extend(entities)
	all_relations.extend(relations)

	# Add to NetworkX graph
	for entity in entities:
	self.graph.add_node(entity.id, **entity.properties)

	for relation in relations:
	self.graph.add_edge(
	relation.subject,
	relation.object,
	predicate=relation.predicate,
	confidence=relation.confidence,
	metadata=relation.metadata,
	)

	# Store entities and relations
	self.entities = {entity.id: entity for entity in all_entities}
	self.relations = all_relations

	# Generate embeddings for entities
	await self._generate_entity_embeddings()

	# Create vector store
	await self._create_vector_store(documents)

	stats = {
	"total_entities": len(all_entities),
	"total_relations": len(all_relations),
	"graph_nodes": self.graph.number_of_nodes(),
	"graph_edges": self.graph.number_of_edges(),
	}

	logger.info(f"Knowledge graph built: {stats}")
	return stats

	async def _generate_entity_embeddings(self):
	"""تولید embeddings برای موجودیت‌ها"""

	logger.info("Generating entity embeddings")

	for entity_id, entity in self.entities.items():
	try:
	# Create description for embedding
	description = f"{entity.name} is a {entity.type}. {entity.properties.get('description', '')}"

	# Generate embedding
	embedding = await self.embeddings.aembed_query(description)
	entity.embeddings = embedding

	except Exception as e:
	logger.error(f"Error generating embedding for {entity_id}: {e}")

	async def _create_vector_store(self, documents: List[str]):
	"""ایجاد vector store از اسناد"""

	logger.info("Creating vector store")

	try:
	# Create documents
	docs = [
	Document(page_content=doc, metadata={"source": f"doc_{i}"})
	for i, doc in enumerate(documents)
	]

	# Split documents
	text_splitter = RecursiveCharacterTextSplitter(
	chunk_size=1000, chunk_overlap=200
	)
	split_docs = text_splitter.split_documents(docs)

	# Create vector store
	self.vector_store = await FAISS.afrom_documents(split_docs, self.embeddings)

	logger.info(f"Vector store created with {len(split_docs)} chunks")

	except Exception as e:
	logger.error(f"Error creating vector store: {e}")

	async def query_knowledge_graph(
	self, query: str, use_rag: bool = True
	) -> Dict[str, Any]:
	"""پرس‌وجو از Knowledge Graph"""

	logger.info(f"Querying knowledge graph: {query}")

	result = {
	"query": query,
	"direct_answers": [],
	"rag_answers": [],
	"graph_paths": [],
	"related_entities": [],
	}

	try:
	# Direct graph query
	if self.graph.number_of_nodes() > 0:
	result["direct_answers"] = await self._query_graph_directly(query)

	# RAG-based query
	if use_rag and self.vector_store:
	result["rag_answers"] = await self._query_with_rag(query)

	# Find graph paths
	result["graph_paths"] = await self._find_relevant_paths(query)

	# Find related entities
	result["related_entities"] = await self._find_related_entities(query)

	except Exception as e:
	logger.error(f"Error querying knowledge graph: {e}")
	result["error"] = str(e)

	return result

	async def _query_graph_directly(self, query: str) -> List[str]:
	"""پرس‌وجو مستقیم از گراف"""

	answers = []

	try:
	# Simple keyword matching
	query_lower = query.lower()

	for node in self.graph.nodes():
	node_data = self.graph.nodes[node]
	if any(
	keyword in str(node_data).lower() for keyword in query_lower.split()
	):
	answers.append(f"Found entity: {node} with properties: {node_data}")

	# Find edges containing query terms
	for edge in self.graph.edges(data=True):
	edge_data = edge[2]
	if any(
	keyword in str(edge_data).lower() for keyword in query_lower.split()
	):
	answers.append(
	f"Found relation: {edge[0]} -> {edge[1]} ({edge_data.get('predicate', 'unknown')})"
	)

	except Exception as e:
	logger.error(f"Error in direct graph query: {e}")

	return answers

	async def _query_with_rag(self, query: str) -> List[str]:
	"""پرس‌وجو با استفاده از RAG"""

	answers = []

	try:
	# Retrieve relevant documents
	docs = await self.vector_store.asimilarity_search(query, k=3)

	# Create RAG chain
	rag_prompt = PromptTemplate(
	input_variables=["context", "question"],
	template="""
	بر اساس اطلاعات زیر به سوال پاسخ دهید:

	اطلاعات مرتبط:
	{context}

	سوال: {question}

	پاسخ:
	""",
	)

	# Combine retrieved documents
	context = "\n\n".join([doc.page_content for doc in docs])

	# Generate answer with Gemini
	prompt = rag_prompt.format(context=context, question=query)
	response = await self.llm.ainvoke(prompt)

	answers.append(response.content)

	except Exception as e:
	logger.error(f"Error in RAG query: {e}")

	return answers

	async def _find_relevant_paths(self, query: str) -> List[Dict[str, Any]]:
	"""یافتن مسیرهای مرتبط در گراف"""

	paths = []

	try:
	query_terms = query.lower().split()

	# Find entities matching query terms
	relevant_nodes = []
	for node in self.graph.nodes():
	node_data = self.graph.nodes[node]
	if any(term in str(node_data).lower() for term in query_terms):
	relevant_nodes.append(node)

	# Find paths between relevant nodes
	if len(relevant_nodes) >= 2:
	for i in range(len(relevant_nodes)):
	for j in range(i + 1, len(relevant_nodes)):
	try:
	path = nx.shortest_path(
	self.graph, relevant_nodes[i], relevant_nodes[j]
	)
	if len(path) > 1:
	paths.append(
	{
	"path": path,
	"length": len(path) - 1,
	"source": relevant_nodes[i],
	"target": relevant_nodes[j],
	}
	)
	except nx.NetworkXNoPath:
	continue

	except Exception as e:
	logger.error(f"Error finding paths: {e}")

	return paths

	async def _find_related_entities(self, query: str) -> List[Dict[str, Any]]:
	"""یافتن موجودیت‌های مرتبط"""

	related = []

	try:
	# Use embeddings for similarity search
	query_embedding = await self.embeddings.aembed_query(query)

	for entity_id, entity in self.entities.items():
	if entity.embeddings:
	# Calculate cosine similarity
	similarity = np.dot(query_embedding, entity.embeddings) / (
	np.linalg.norm(query_embedding)
	* np.linalg.norm(entity.embeddings)
	)

	if similarity > 0.7: # Threshold
	related.append(
	{
	"entity_id": entity_id,
	"name": entity.name,
	"type": entity.type,
	"similarity": float(similarity),
	"properties": entity.properties,
	}
	)

	# Sort by similarity
	related.sort(key=lambda x: x["similarity"], reverse=True)

	except Exception as e:
	logger.error(f"Error finding related entities: {e}")

	return related[:5] # Return top 5

	async def create_rag_agent(self) -> Any:
	"""ایجاد RAG Agent با Gemini"""

	try:
	# Define tools
	tools = [
	Tool(
	name="knowledge_graph_search",
	description="جستجو در Knowledge Graph برای یافتن اطلاعات",
	func=lambda query: str(self.query_knowledge_graph(query)),
	),
	Tool(
	name="web_search",
	description="جستجو در اینترنت برای اطلاعات جدید",
	func=self.search.run,
	),
	]

	# Create agent
	agent = initialize_agent(
	tools=tools,
	llm=self.llm,
	agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
	memory=self.memory,
	verbose=True,
	handle_parsing_errors=True,
	)

	logger.info("RAG Agent created successfully")
	return agent

	except Exception as e:
	logger.error(f"Error creating RAG agent: {e}")
	return None

	def visualize_knowledge_graph(
	self, max_nodes: int = 50, save_path: str = "kg_visualization.png"
	):
	"""تجسم Knowledge Graph"""

	try:
	plt.figure(figsize=(20, 16))

	# Create subgraph if too many nodes
	if self.graph.number_of_nodes() > max_nodes:
	# Get most connected nodes
	centrality = nx.degree_centrality(self.graph)
	top_nodes = sorted(
	centrality.items(), key=lambda x: x[1], reverse=True
	)[:max_nodes]
	subgraph = self.graph.subgraph([node for node, _ in top_nodes])
	else:
	subgraph = self.graph

	# Layout
	pos = nx.spring_layout(subgraph, k=3, iterations=50)

	# Draw nodes
	nx.draw_networkx_nodes(
	subgraph, pos, node_color="lightblue", node_size=1000, alpha=0.7
	)

	# Draw edges
	nx.draw_networkx_edges(
	subgraph, pos, edge_color="gray", arrows=True, arrowsize=20, alpha=0.6
	)

	# Draw labels
	nx.draw_networkx_labels(subgraph, pos, font_size=8, font_weight="bold")

	plt.title("Knowledge Graph Visualization", fontsize=16, fontweight="bold")
	plt.axis("off")
	plt.tight_layout()
	plt.savefig(save_path, dpi=300, bbox_inches="tight")
	plt.show()

	logger.info(f"Knowledge graph visualization saved to {save_path}")

	except Exception as e:
	logger.error(f"Error visualizing knowledge graph: {e}")

	def save_knowledge_graph(self, filepath: str):
	"""ذخیره Knowledge Graph"""

	try:
	kg_data = {
	"entities": {
	k: {
	"id": v.id,
	"name": v.name,
	"type": v.type,
	"properties": v.properties,
	"embeddings": v.embeddings,
	}
	for k, v in self.entities.items()
	},
	"relations": [
	{
	"subject": r.subject,
	"predicate": r.predicate,
	"object": r.object,
	"confidence": r.confidence,
	"metadata": r.metadata,
	}
	for r in self.relations
	],
	"graph_data": {
	"nodes": list(self.graph.nodes(data=True)),
	"edges": list(self.graph.edges(data=True)),
	},
	"metadata": {
	"created_at": datetime.now().isoformat(),
	"total_entities": len(self.entities),
	"total_relations": len(self.relations),
	},
	}

	with open(filepath, "w", encoding="utf-8") as f:
	json.dump(kg_data, f, ensure_ascii=False, indent=2)

	logger.info(f"Knowledge graph saved to {filepath}")

	except Exception as e:
	logger.error(f"Error saving knowledge graph: {e}")

	def load_knowledge_graph(self, filepath: str):
	"""بارگذاری Knowledge Graph"""

	try:
	with open(filepath, "r", encoding="utf-8") as f:
	kg_data = json.load(f)

	# Load entities
	self.entities = {}
	for k, v in kg_data["entities"].items():
	entity = KGEntity(
	id=v["id"],
	name=v["name"],
	type=v["type"],
	properties=v["properties"],
	embeddings=v.get("embeddings"),
	)
	self.entities[k] = entity

	# Load relations
	self.relations = []
	for r in kg_data["relations"]:
	relation = KGRelation(
	subject=r["subject"],
	predicate=r["predicate"],
	object=r["object"],
	confidence=r["confidence"],
	metadata=r["metadata"],
	)
	self.relations.append(relation)

	# Rebuild NetworkX graph
	self.graph.clear()
	for node, data in kg_data["graph_data"]["nodes"]:
	self.graph.add_node(node, **data)

	for edge in kg_data["graph_data"]["edges"]:
	self.graph.add_edge(edge[0], edge[1], **edge[2])

	logger.info(f"Knowledge graph loaded from {filepath}")

	except Exception as e:
	logger.error(f"Error loading knowledge graph: {e}")


	# Example usage and testing functions
	async def run_experimental_examples():
	"""اجرای مثال‌های تجربی"""

	# Initialize system
	api_key = os.getenv("GEMINI_API_KEY")
	if not api_key:
	logger.error("GEMINI_API_KEY not found in environment variables")
	return

	kg_system = GeminiKnowledgeGraphSystem(api_key)

	# Sample documents
	sample_documents = [
	"""
	آلبرت انیشتین یک فیزیکدان نظری آلمانی بود که نظریه نسبیت را توسعه داد.
	او در سال 1879 در اولم، آلمان متولد شد و در سال 1955 در پرینستون، نیوجرسی درگذشت.
	انیشتین جایزه نوبل فیزیک را در سال 1921 دریافت کرد.
	""",
	"""
	نظریه نسبیت انیشتین شامل دو بخش اصلی است: نسبیت خاص و نسبیت عام.
	نسبیت خاص در سال 1905 منتشر شد و شامل رابطه E=mc² است.
	نسبیت عام در سال 1915 منتشر شد و گرانش را به عنوان انحنای فضا-زمان توصیف می‌کند.
	""",
	"""
	دانشگاه پرینستون یکی از قدیمی‌ترین و معتبرترین دانشگاه‌های آمریکا است.
	انیشتین از سال 1933 تا زمان مرگش در سال 1955 در موسسه مطالعات پیشرفته پرینستون کار کرد.
	این دانشگاه در نیوجرسی واقع شده است.
	""",
	]

	try:
	# Build knowledge graph
	print("🔨 Building Knowledge Graph...")
	stats = await kg_system.build_knowledge_graph(sample_documents)
	print(f"✅ Knowledge Graph built: {stats}")

	# Query examples
	queries = [
	"انیشتین کی بود؟",
	"نظریه نسبیت چیست؟",
	"پرینستون کجاست؟",
	"ارتباط انیشتین با پرینستون چیست؟",
	]

	for query in queries:
	print(f"\n🔍 Query: {query}")
	result = await kg_system.query_knowledge_graph(query)

	print("📝 Direct Answers:")
	for answer in result.get("direct_answers", []):
	print(f" - {answer}")

	print("🤖 RAG Answers:")
	for answer in result.get("rag_answers", []):
	print(f" - {answer}")

	print("🛤️ Graph Paths:")
	for path in result.get("graph_paths", []):
	print(f" - {' -> '.join(path['path'])}")

	print("🔗 Related Entities:")
	for entity in result.get("related_entities", []):
	print(
	f" - {entity['name']} ({entity['type']}) - Similarity: {entity['similarity']:.3f}"
	)

	# Create RAG Agent
	print("\n🤖 Creating RAG Agent...")
	agent = await kg_system.create_rag_agent()

	if agent:
	print("✅ RAG Agent created successfully")

	# Test agent
	agent_queries = [
	"درباره انیشتین چه می‌دانید؟",
	"آیا اطلاعات جدیدی درباره نظریه نسبیت وجود دارد؟",
	]

	for query in agent_queries:
	print(f"\n🤖 Agent Query: {query}")
	try:
	response = await agent.arun(query)
	print(f"🤖 Agent Response: {response}")
	except Exception as e:
	print(f"❌ Agent Error: {e}")

	# Visualize knowledge graph
	print("\n📊 Visualizing Knowledge Graph...")
	kg_system.visualize_knowledge_graph()

	# Save knowledge graph
	print("\n💾 Saving Knowledge Graph...")
	kg_system.save_knowledge_graph("experimental_kg.json")

	print("\n🎉 Experimental examples completed successfully!")

	except Exception as e:
	logger.error(f"Error in experimental examples: {e}")


	if __name__ == "__main__":
	# Run experimental examples
	asyncio.run(run_experimental_examples())