mirror of
https://github.com/Mintplex-Labs/langchain-python.git
synced 2026-07-18 18:34:27 -04:00
maspotts
95c9aa1ccb
Copies `GraphIndexCreator.from_text()` to make an async version called `GraphIndexCreator.afrom_text()`. This is (should be) a trivial change: it just adds a copy of `GraphIndexCreator.from_text()` which is async and awaits a call to `chain.apredict()` instead of `chain.predict()`. There is no unit test for GraphIndexCreator, and I did not create one, but this code works for me locally. @agola11 @hwchase17
43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
"""Graph Index Creator."""
|
|
from typing import Optional, Type
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from langchain.base_language import BaseLanguageModel
|
|
from langchain.chains.llm import LLMChain
|
|
from langchain.graphs.networkx_graph import NetworkxEntityGraph, parse_triples
|
|
from langchain.indexes.prompts.knowledge_triplet_extraction import (
|
|
KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT,
|
|
)
|
|
|
|
|
|
class GraphIndexCreator(BaseModel):
|
|
"""Functionality to create graph index."""
|
|
|
|
llm: Optional[BaseLanguageModel] = None
|
|
graph_type: Type[NetworkxEntityGraph] = NetworkxEntityGraph
|
|
|
|
def from_text(self, text: str) -> NetworkxEntityGraph:
|
|
"""Create graph index from text."""
|
|
if self.llm is None:
|
|
raise ValueError("llm should not be None")
|
|
graph = self.graph_type()
|
|
chain = LLMChain(llm=self.llm, prompt=KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT)
|
|
output = chain.predict(text=text)
|
|
knowledge = parse_triples(output)
|
|
for triple in knowledge:
|
|
graph.add_triple(triple)
|
|
return graph
|
|
|
|
async def afrom_text(self, text: str) -> NetworkxEntityGraph:
|
|
"""Create graph index from text asynchronously."""
|
|
if self.llm is None:
|
|
raise ValueError("llm should not be None")
|
|
graph = self.graph_type()
|
|
chain = LLMChain(llm=self.llm, prompt=KNOWLEDGE_TRIPLE_EXTRACTION_PROMPT)
|
|
output = await chain.apredict(text=text)
|
|
knowledge = parse_triples(output)
|
|
for triple in knowledge:
|
|
graph.add_triple(triple)
|
|
return graph
|