from dotenv import load_dotenv
import dspy
from modaic import PrecompiledProgram, PrecompiledConfig

load_dotenv()

class CypherFromText(dspy.Signature):
    """Instructions:
    Create a Cypher MERGE statement to model all entities and relationships found in the text following these guidelines:
    - Refer to the provided schema and use existing or similar nodes, properties or relationships before creating new ones.
    - Use generic categories for node and relationship labels."""
    text = dspy.InputField(
        desc="Text to model using nodes, properties and relationships."
    )
    neo4j_schema = dspy.InputField(
        desc="Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS."
    )
    statement = dspy.OutputField(
        desc="Cypher statement to merge nodes and relationships found in the text."
    )

class GenerateCypherConfig(PrecompiledConfig):
    neo4j_schema: list[str] = []
    model: str = "gpt-4"
    max_tokens: int = 1024
    

class GenerateCypher(PrecompiledProgram):
    config: GenerateCypherConfig

    def _init_(self, config: GenerateCypherConfig, **kwargs):
        super()._init_(**kwargs)
        self.lm = dspy.LM(
            model=config.model,
            max_tokens=config.max_tokens,
        )
        self.generate_cypher = dspy.ChainOfThought(CypherFromText)
    
    def forward(self, text: str, neo4j_schema: list[str]):
        return self.generate_cypher(text=text, neo4j_schema=neo4j_schema)


if __name__ == "__main__":
    generate_cypher = GenerateCypher(GenerateCypherConfig())
    generate_cypher.push_to_hub("farouk1/text-to-cypher", with_code=True, tag="v0.0.1", commit_message="init")