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

fix model_tokens not being used for ollama

Browse Source 
I am passing in the explicit model_tokens from user config as the default_token so it will correctly fallback to the users setting if the model is not found
This commit is contained in:
Jamie Beck 2024-08-26 11:30:09 -04:00 committed by GitHub
parent 88e76ceedb
commit 35b994a8cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
scrapegraphai/graphs/abstract_graph.py
View File

@ -174,8 +174,9 @@ class AbstractGraph(ABC):
elif "ollama" in llm_params["model"]:
model_name = llm_params["model"].split("ollama/")[-1]
token_key = model_name if "model_tokens" not in llm_params else llm_params["model_tokens"]
return handle_model(model_name, "ollama", token_key)
token_key = model_name if "model_tokens" not in llm_params else None
explicit_model_tokens = 8192 if "model_tokens" not in llm_params else llm_params["model_tokens"]
return handle_model(model_name, "ollama", token_key, explicit_model_tokens)
elif "claude-3-" in llm_params["model"]:
return handle_model(llm_params["model"], "anthropic", "claude3")
@ -271,4 +272,4 @@ class AbstractGraph(ABC):
def run(self) -> str:
"""
Abstract method to execute the graph and return the result.
"""
"""