a dgE4@sddlmZddlmZmZmZmZmZmZm Z ddl Z ddl m Z ddl mZmZmZddlmZGdddeeZeed d d ZGd d d ZgdZgdZGdddZGdddZdS))Enum)AnyDictListOptionalSequenceTupleUnionN) get_from_env) GraphDocumentNode Relationship)Documentc@seZdZdZdZdZdS) TypeOptionfactsentities sentimentN)__name__ __module__ __qualname__FACTSZENTITIESZ SENTIMENTrr/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_experimental/graph_transformers/diffbot.pyr sr)sreturncCsB|}|s|S|d}dd|ddD}d|g|S)z.Formats a string to be used as a property key.rcSsg|] }|qSr) capitalize).0wordrrr z'format_property_key..N)splitlowerjoin)rwordsZ first_wordZcapitalized_wordsrrrformat_property_keys  r&c@sZeZdZdZddddZeeeefefe ee fddddZ e e dd d ZdS) NodesLista List of nodes with associated properties. Attributes: nodes (Dict[Tuple, Any]): Stores nodes as keys and their properties as values. Each key is a tuple where the first element is the node ID and the second is the node type. NrcCs t|_dS)N)dictnodesselfrrr__init__$szNodesList.__init__)node propertiesrcCs*||jvr||j|<n|j||dS)a Adds or updates node properties. If the node does not exist in the list, it's added along with its properties. If the node already exists, its properties are updated with the new values. Args: node (Tuple): A tuple containing the node ID and node type. properties (Dict): A dictionary of properties to add or update for the node. N)r*update)r,r.r/rrradd_node_property's  zNodesList.add_node_propertycsfddjD}|S)z Returns the nodes as a list of Node objects. Each Node object will have its ID, type, and properties populated. Returns: List[Node]: A list of Node objects. cs(g|] }t|d|dj|dqS)rr )idtyper/)r r*)rkeyr+rrrBsz.NodesList.return_node_list..)r*)r,r*rr+rreturn_node_list9s zNodesList.return_node_list)rrr__doc__r-rr strintrrr1rr r5rrrrr's  r')DateNumberz Job titlezCause of deathzOrganization typezAcademic title))Z HEADQUARTERSZORGANIZATION_LOCATIONS)Z RESIDENCEPERSON_LOCATION)ZALL_PERSON_LOCATIONSr;)ZCHILDZ HAS_CHILD)ZPARENTZ HAS_PARENT)Z CUSTOMERSZ HAS_CUSTOMER)Z SKILLED_ATZ INTERESTED_INc@s.eZdZdZddddZeedddZdS) SimplifiedSchemazSimplified schema mapping. Attributes: schema (Dict): A dictionary containing the mapping to simplified schema types. Nr(cCs(t|_tD]}|d|j|d<q dS)z?Initializes the schema dictionary based on the predefined list.r rN)r)schemaschema_mapping)r,rowrrrr-fszSimplifiedSchema.__init__)r3rcCs(z |j|WSty"|YS0dS)a5 Retrieves the simplified schema type for a given original type. Args: type (str): The original schema type to find the simplified type for. Returns: str: The simplified schema type if it exists; otherwise, returns the original type. N)r=KeyError)r,r3rrrget_typels   zSimplifiedSchema.get_type)rrrr6r-r7rArrrrr<_sr<c @seZdZdZdddddejgfddeeee e e e ee dddd Z ee ee fd d d Ze ee feed ddZeee edddZdS)DiffbotGraphTransformeraTransform documents into graph documents using Diffbot NLP API. A graph document transformation system takes a sequence of Documents and returns a sequence of Graph Documents. Example: .. code-block:: python from langchain_experimental.graph_transformers import DiffbotGraphTransformer from langchain_core.documents import Document diffbot_api_key = "DIFFBOT_API_KEY" diffbot_nlp = DiffbotGraphTransformer(diffbot_api_key=diffbot_api_key) document = Document(page_content="Mike Tunge is the CEO of Diffbot.") graph_documents = diffbot_nlp.convert_to_graph_documents([document]) Ngffffff?TF)include_confidence)diffbot_api_keyfact_confidence_thresholdinclude_qualifiersinclude_evidencesimplified_schema extract_typesrCrcCsP|p tdd|_||_||_||_||_d|_|r:t|_|sFtd||_ dS)a  Initialize the graph transformer with various options. Args: diffbot_api_key (str): The API key for Diffbot's NLP services. fact_confidence_threshold (float): Minimum confidence level for facts to be included. include_qualifiers (bool): Whether to include qualifiers in the relationships. include_evidence (bool): Whether to include evidence for the relationships. simplified_schema (bool): Whether to use a simplified schema for relationships. extract_types (List[TypeOption]): A list of data types to extract. Facts, entities, and sentiment are supported. By default, the option is set to facts. A fact represents a combination of source and target nodes with a relationship type. include_confidence (bool): Whether to include confidence scores on nodes and rels rDZDIFFBOT_API_KEYNzZ`extract_types` cannot be an empty array. Allowed values are 'facts', 'entities', or both.) r rDfact_threshold_confidencerFrGrCrHr< ValueErrorrI)r,rDrErFrGrHrIrCrrrr-s"z DiffbotGraphTransformer.__init__)textrcCsJ|dd}d|j}d}d|d|d|jd}tj||d }|S) z Make an API request to the Diffbot NLP endpoint. Args: text (str): The text to be processed. Returns: Dict[str, Any]: The JSON response from the API. en)contentlang,znl.diffbot.comzhttps://z /v1/?fields=z&token=z &language=en)data)r$rIrDrequestspostjson)r,rLpayloadZFIELDSZHOSTurlresultrrr nlp_requests  z#DiffbotGraphTransformer.nlp_request)rUdocumentrcCs2d|vs|ds.d|vs |ds.tgg|dSt}d|vr|dr|dD]}|dsZqL|drn|ddn|d}|ddd}|d}|||fd|i|dd ur|||fd|di|jrL|||fd |d iqLt}d|vr |dr |dD] }|d |jkr,q|d ds>q|d dr\|d ddn |d d}|d ddd}|d d}t||d } |||fd|i|d dr|d ddn |d d} |d ddd} |d d} | t vr&|||ft |dd| int| | d } || | fd| i|dd dd }|j rr|j |}t}dd|dDd}|jr|d|i|jr|d |d i|jr|dr|dD]&}t |dd}|d d||<qt| | ||d}||qt|||dS)a3 Transform the Diffbot NLP response into a GraphDocument. Args: payload (Dict[str, Any]): The JSON response from Diffbot's NLP API. document (Document): The original document. Returns: GraphDocument: The transformed document as a graph. rr)r* relationshipssourceZallTypesZallUrisrnamerN confidencevalueentity)r2r3property _cSsg|] }|dqS)Zpassager)relrrrrBsz.ZevidenceZ qualifiers)r[targetr3r/)r r'rr1getrClistrJr FACT_TO_PROPERTY_TYPEr&replaceupperrHrAr)rGr0rFr appendr5)r,rUrYZ nodes_listrecordZ source_idZ source_labelZ source_namerZZ source_nodeZ target_idZ target_labelZ target_nameZ target_nodeZrel_typeZrel_propertiesZrelationship_evidencer`Zprop_keyZ relationshiprrrprocess_responses                 z(DiffbotGraphTransformer.process_response) documentsrcCs4g}|D]&}||j}|||}||q|S)aConvert a sequence of documents into graph documents. Args: documents (Sequence[Document]): The original documents. kwargs: Additional keyword arguments. Returns: Sequence[GraphDocument]: The transformed documents as graphs. )rXZ page_contentrlrj)r,rmresultsrYZ raw_resultsZgraph_documentrrrconvert_to_graph_documents\s    z2DiffbotGraphTransformer.convert_to_graph_documents)rrrr6rrrr7floatboolrr-rrrXrr rlrrorrrrrB}s4  4  rB)enumrtypingrrrrrrr rRZlangchain.utilsr Z)langchain_community.graphs.graph_documentr r r Zlangchain_core.documentsrr7rr&r'rgr>r<rBrrrrs $   /