from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader("/azdev/paulgraham/").load_data()

import nest_asyncio
from typing import Literal
from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.core.indices.property_graph import SchemaLLMPathExtractor

nest_asyncio.apply()

entities = Literal["PERSON", "PLACE", "ORGANIZATION"]
relations = Literal["HAS", "PART_OF", "WORKED_ON", "WORKED_WITH", "WORKED_AT"]

validation_schema = {
    "PERSON": ["HAS", "PART_OF", "WORKED_ON", "WORKED_WITH", "WORKED_AT"],
    "PLACE": ["HAS", "PART_OF", "WORKED_AT"],
    "ORGANIZATION": ["HAS", "PART_OF", "WORKED_WITH"],
}

kg_extractor = SchemaLLMPathExtractor(
    llm=AzureOpenAI(
        model="gpt-4",
        deployment_name="",
        api_key='',
        azure_endpoint="",
        api_version="2024-02-01",
    ),
    possible_entities=entities,
    possible_relations=relations,
    kg_validation_schema=validation_schema,
    strict=True
)

from llama_index.graph_stores.neo4j import Neo4jPGStore

graph_store = Neo4jPGStore(
    username="neo4j",
    password="",
    url="bolt://localhost:7687",
)

import os
import asyncio
from openai import AzureOpenAI
from llama_index.core import PropertyGraphIndex

os.environ["AZURE_OPENAI_API_KEY"] = ""
os.environ["AZURE_OPENAI_ENDPOINT"] = ""

client = AzureOpenAI(
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    api_version="2024-02-01",
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
)

deployment_name = 'embedding'

def azure_openai_embedding(text):
    try:
        response = client.embeddings.create(
            model=deployment_name,
            input=text
        )
        return response['data'][0]['embedding']
    except Exception as e:
        print(f"Error getting embedding: {e}")
        return None

class AzureOpenAIEmbedding:
    def __init__(self, client):
        self.client = client

    def embed(self, text):
        return azure_openai_embedding(text)

    async def aget_text_embedding_batch(self, texts, **kwargs):
        try:
            loop = asyncio.get_event_loop()
            tasks = [loop.run_in_executor(None, azure_openai_embedding, text) for text in texts]
            responses = await asyncio.gather(*tasks)
            return [response for response in responses if response is not None]
        except Exception as e:
            print(f"Error getting batch embeddings: {e}")
            return None

embed_model = AzureOpenAIEmbedding(client)

try:
    index = PropertyGraphIndex.from_documents(
        documents,
        kg_extractors=[kg_extractor],
        embed_model=embed_model,
        property_graph_store=graph_store,
    )
except Exception as e:
    print(f"Error creating PropertyGraphIndex: {e}")