+#hEddlZddlmZmZmZmZddlZddlmZm Z m Z m Z m Z ddl Z ddlmZ ddlmZddlmZmZGdd Zd ZGd d e j0Zy#e$r edwxYw) N)AnyListOptionalUnion)APITimeoutErrorInternalServerErrorOpenAIRateLimitErrorUnprocessableEntityError)dotdict) GraphDatabase) AuthErrorServiceUnavailablezHPlease install the neo4j package by running `pip install dspy-ai[neo4j]`cteZdZdedefdZej ejee e e fdde fdZ y) Embedderprovidermodelc|dk(r9tjd}|s tdt|_||_yy)NopenaiOPENAI_API_KEYz/Environment variable OPENAI_API_KEY must be set)osgetenv ValueErrorr clientr)selfrrapi_keys S/var/www/html/sandstorm/venv/lib/python3.12/site-packages/dspy/retrieve/neo4j_rm.py__init__zEmbedder.__init__s= x ii 01G !RSS (DKDJ )max_timereturnc|jjj||j}|jDcgc]}|j c}Scc}w)N)inputr)r embeddingscreaterdata embedding)rqueriesr(results r__call__zEmbedder.__call__%sEKK**11tzz1R /8~~>V  >>>sAN)__name__ __module__ __qualname__strrbackoff on_exceptionexporrr r rr+rrrrsZSW    $   ?3? ?rrzJCALL db.index.vector.queryNodes($index, $k, $embedding) YIELD node, score c eZdZdZ d dedededededef fd Zd eeeefde ed e jfd Z xZ S) Neo4jRMa Implements a retriever that utilizes Neo4j for retrieving passages. This class manages a connection to a Neo4j database using official Neo4j Python drivers and requires the database credentials (username, password, URI, and optionally the database name) to be set as environment variables. Additionally, it utilizes an embedding provider (defaulting to OpenAI's services) to compute query embeddings, which are then used to find the most relevant nodes in the Neo4j graph based on the specified node property or custom retrieval query. Returns a list of passages in the form of `dspy.Prediction` objects. Args: index_name (str): The name of the vector index in the Neo4j database to query against. text_node_property (Optional[str]): The property of the node containing the text. Required if `retrieval_query` is not set. k (Optional[int]): The default number of top passages to retrieve. Defaults to 5. retrieval_query (Optional[str]): Custom Cypher query for retrieving passages. Required if `text_node_property` is not set. embedding_provider (str): The provider of the embedding service. Defaults to "openai". embedding_model (str): The model identifier for generating embeddings. Defaults to "text-embedding-ada-002". Examples: Below is a code snippet showcasing how to initialize Neo4jRM with environment variables for the database connection and OpenAI as the embedding provider: ```python import os import dspy import openai os.environ["NEO4J_URI"] = "bolt://localhost:7687" os.environ["NEO4J_USERNAME"] = "neo4j" os.environ["NEO4J_PASSWORD"] = "password" os.environ["OPENAI_API_KEY"] = "sk-" neo4j_retriever = Neo4jRM( index_name="myIndex", text_node_property="text", k=10, embedding_provider="openai", embedding_model="text-embedding-ada-002", ) dspy.settings.configure(rm=neo4j_retriever) ``` In this example, `Neo4jRM` is configured to retrieve nodes based on the "text" property from an index named "myIndex", using embeddings computed by OpenAI's "text-embedding-ada-002" model. index_nametext_node_propertykretrieval_queryembedding_providerembedding_modelct||||_tjd|_tjd|_tjd|_tjdd|_||_ ||_ ||_ |j s td|j s td|js td |js|js td  tj|j|j |j f |_|jjt'|| |_y#t t"f$r}t%d |d}~wwxYw)N)r8NEO4J_USERNAMENEO4J_PASSWORD NEO4J_URINEO4J_DATABASEneo4jz/Environment variable NEO4J_USERNAME must be setz/Environment variable NEO4J_PASSWORD must be setz*Environment variable NEO4J_URI must be setzKEither `text_node_property` or `retrieval_query` parameters must be defined)authz#Failed to connect to Neo4j database)rr)superrr6rrusernamepassworduridatabaser8r9r7rr driververify_connectivityrrConnectionErrorrembedder) rr6r7r8r9r:r;e __class__s rrzNeo4jRM.__init__fsF 1$ "23  "23 99[) "2G< ."4}}NO O}}NO OxxIJ J&&t/C/Cjk k P'..txxt}}dmm>\]DK KK + + -!*z!Neo4jRM.forward..s !G*rT)keyreverserU) isinstancelistrKr9r7rH execute_queryDEFAULT_INDEX_QUERYr6r8rGextendr sorted) rrNr8 query_vectorscontentsr9vectorrecords_rsorted_passagesels rforwardzNeo4jRM.forwards *D1 01  &67 ..iLAXAX@YYh2i# tF KK55#o5$tQ[$&&Q--6MGQ OOjqref+qy1I)JUVW^U_`r s  t! $  AK  )88"9 88 s 9s (D 2 D)NNrztext-embedding-ada-002)r,r-r.__doc__r/intrrrrdspy Predictionri __classcell__)rMs@rr5r57s,b#'#"*7$U$U $U  $U  $U  $U$UL9c49n(=9(3-9TXTcTc9rr5)rtypingrrrrr0rrrr r r rm dsp.utilsr rAr neo4j.exceptionsrr ImportErrorrr^Retriever5r3rrrusw --  #??0cg9dmmg9C R s AA,