apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-29 18:36:22 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Fixed HarmCategory API

Browse source
This commit is contained in:
JackColquitt 2024-10-10 11:37:21 -07:00
parent 5850b6c136
commit 063acc6a21
1 changed files with 4 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

6
trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py
View file

@ -5,6 +5,7 @@ Input is prompt, output is response.
"""
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold
from prometheus_client import Histogram
from .... schema import TextCompletionRequest, TextCompletionResponse, Error
@ -80,8 +81,9 @@ class Processor(ConsumerProducer):
HarmCategory.HARM_CATEGORY_HATE_SPEECH: block_level,
HarmCategory.HARM_CATEGORY_HARASSMENT: block_level,
HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: block_level,
HarmCategory.HARM_CATEGORY_DANGEROUS: block_level,
HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level,
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: block_level,
# There is a documentation conflict on whether or not CIVIC_INTEGRITY is a valid category
# HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level,
}
self.llm = genai.GenerativeModel(