CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-07-12 21:01:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

complete get_state function

Browse Source
This commit is contained in:
VinciGit00 2024-04-25 15:08:42 +02:00
parent 1b004d82a1
commit 75255bb4af
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
scrapegraphai/graphs/abstract_graph.py
View File

@ -41,7 +41,7 @@ class AbstractGraph(ABC):
try:
self.model_token = models_tokens["openai"][llm_params["model"]]
except KeyError:
raise ValueError("Model not supported")
raise KeyError("Model not supported")
return OpenAI(llm_params)
elif "azure" in llm_params["model"]:
@ -50,14 +50,14 @@ class AbstractGraph(ABC):
try:
self.model_token = models_tokens["azure"][llm_params["model"]]
except KeyError:
raise ValueError("Model not supported")
raise KeyError("Model not supported")
return AzureOpenAI(llm_params)
elif "gemini" in llm_params["model"]:
try:
self.model_token = models_tokens["gemini"][llm_params["model"]]
except KeyError:
raise ValueError("Model not supported")
raise KeyError("Model not supported")
return Gemini(llm_params)
elif "ollama" in llm_params["model"]:
@ -70,19 +70,27 @@ class AbstractGraph(ABC):
try:
self.model_token = models_tokens["ollama"][llm_params["model"]]
except KeyError:
raise ValueError("Model not supported")
raise KeyError("Model not supported")
return Ollama(llm_params)
elif "hugging_face" in llm_params["model"]:
try:
self.model_token = models_tokens["hugging_face"][llm_params["model"]]
except KeyError:
raise ValueError("Model not supported")
raise KeyError("Model not supported")
return HuggingFace(llm_params)
else:
raise ValueError(
"Model provided by the configuration not supported")
def get_state(self, key=None) -> dict:
"""""
Obtain the current state
"""
if key is not None:
return self.final_state[key]
return self.final_state
def get_execution_info(self):
"""
Returns the execution information of the graph.