Skip to content

Commit 47cfdc4

Browse files
seanzhangkx8seanzhangkx8
authored
Add support to Gemini API key for GenAI module (#347)
1 parent ca2646d commit 47cfdc4

4 files changed

Lines changed: 137 additions & 35 deletions

File tree

‎convokit/genai/example/example.ipynb‎

Lines changed: 100 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,19 @@
2222
},
2323
{
2424
"cell_type": "code",
25-
"execution_count": 23,
25+
"execution_count": 1,
2626
"id": "be24454d",
2727
"metadata": {},
2828
"outputs": [],
2929
"source": [
30-
"from convokit.genai.genai_config import GenAIConfigManager"
30+
"from convokit.genai.genai_config import GenAIConfigManager\n",
31+
"config = GenAIConfigManager()\n",
32+
"from convokit.genai import get_llm_client"
3133
]
3234
},
3335
{
3436
"cell_type": "code",
35-
"execution_count": null,
37+
"execution_count": 2,
3638
"id": "5666dc72",
3739
"metadata": {},
3840
"outputs": [
@@ -45,7 +47,6 @@
4547
}
4648
],
4749
"source": [
48-
"config = GenAIConfigManager()\n",
4950
"config.set_api_key(\"gpt\", \"YOUR API KEY\")\n",
5051
"print(f\"Successfully set OpenAI API key in config.\")"
5152
]
@@ -62,13 +63,11 @@
6263
},
6364
{
6465
"cell_type": "code",
65-
"execution_count": 25,
66+
"execution_count": null,
6667
"id": "220c33fe",
6768
"metadata": {},
6869
"outputs": [],
6970
"source": [
70-
"from convokit.genai import get_llm_client\n",
71-
"\n",
7271
"MODEL_PROVIDER = \"gpt\"\n",
7372
"client = get_llm_client(MODEL_PROVIDER, config)"
7473
]
@@ -96,6 +95,95 @@
9695
"print(\"Latency:\", response.latency)"
9796
]
9897
},
98+
{
99+
"cell_type": "markdown",
100+
"id": "417c22b9",
101+
"metadata": {},
102+
"source": [
103+
"# Setup ConvoKit GenAI with Google Gemini Through Gemini Developer API"
104+
]
105+
},
106+
{
107+
"cell_type": "markdown",
108+
"id": "94aede45",
109+
"metadata": {},
110+
"source": [
111+
"Checkout Google Gemini Developer API: https://ai.google.dev/gemini-api/docs/"
112+
]
113+
},
114+
{
115+
"cell_type": "code",
116+
"execution_count": null,
117+
"id": "0a4c8e31",
118+
"metadata": {},
119+
"outputs": [],
120+
"source": [
121+
"MODEL_PROVIDER = \"gemini\"\n",
122+
"MODEL = \"gemini-2.5-flash\"\n",
123+
"config.set_api_key(\"gemini\", \"YOUR API KEY\")"
124+
]
125+
},
126+
{
127+
"cell_type": "code",
128+
"execution_count": 8,
129+
"id": "88eef33f",
130+
"metadata": {},
131+
"outputs": [
132+
{
133+
"name": "stdout",
134+
"output_type": "stream",
135+
"text": [
136+
"Text: ConvoKit is a powerful Python library designed for the **systematic analysis of conversational data**. It provides a unified framework to represent, transform, and analyze conversations, moving beyond simple text analysis to focus on the rich structure and dynamics inherent in human (and human-like) interactions.\n",
137+
"\n",
138+
"Here's a breakdown of what you can do with ConvoKit:\n",
139+
"\n",
140+
"1. **Represent Conversational Data:**\n",
141+
" * **Standardized Structure:** ConvoKit allows you to load diverse conversational datasets (chat logs, forum posts, emails, transcripts, etc.) into a standardized `Corpus` object, which consists of `Utterance` and `Speaker` objects. This makes it easy to work with different datasets consistently.\n",
142+
" * **Metadata Storage:** You can attach arbitrary metadata to conversations, utterances, and speakers (e.g., speaker demographics, conversation topics, timestamps, sentiment scores).\n",
143+
"\n",
144+
"2. **Preprocess and Enrich Conversational Data:**\n",
145+
" * **Tokenization & Linguistic Features:** Integrate with spaCy to perform tokenization, part-of-speech tagging, named entity recognition, and dependency parsing on utterances.\n",
146+
" * **Sentiment Analysis:** Apply pre-trained sentiment models to utterances to understand emotional tone.\n",
147+
" * **Dialog Act Classification:** Use models to classify the communicative function of an utterance (e.g., question, statement, agreement, disagreement).\n",
148+
" * **Speaker Attributes:** Extract features related to speakers, such as their average utterance length, lexical diversity, or frequency of certain dialog acts.\n",
149+
"\n",
150+
"3. **Analyze Conversational Dynamics and Structure:**\n",
151+
" * **Turn-Taking Patterns:** Investigate how speakers take turns, who initiates conversations, and who responds.\n",
152+
" * **Interaction Networks:** Build and analyze networks of who talks to whom, identifying central speakers or conversational hubs.\n",
153+
" * **Topic Modeling:** Discover latent topics within conversations and track how topics shift over time or across different speakers.\n",
154+
" * **Power and Dominance:** Analyze linguistic cues and interaction patterns that might indicate power dynamics or dominance within a conversation (e.g., interruption patterns, question-asking behavior).\n",
155+
" * **Coherence and Cohesion:** Study how utterances relate to each other to form coherent conversations.\n",
156+
"\n",
157+
"4. **Extract Features for Machine Learning:**\n",
158+
" * **Linguistic Features:** Extract features like word counts, TF-IDF, n-grams, lexical diversity, specific part-of-speech frequencies.\n",
159+
" * **Social Features:** Features related to speaker roles, interaction frequency, or network centrality.\n",
160+
" * **Structural Features:** Features related to turn-taking, utterance position in a conversation, or response time.\n",
161+
" * These features can then be used to train models for tasks like speaker classification, conversation summarization, or predicting conversation outcomes.\n",
162+
"\n",
163+
"5. **Conduct Comparative Studies:**\n",
164+
" * **Compare Groups:** Analyze differences in communication styles, sentiment, or interaction patterns between different groups of speakers (e.g., male vs. female, expert vs. novice, different customer segments).\n",
165+
" * **Compare Contexts:** Study how conversations differ across various contexts (e.g., formal vs. informal, online vs. offline, different platforms).\n",
166+
"\n",
167+
"6. **Specific Use Cases:**\n",
168+
"\n",
169+
" * **Social Science Research:** Studying group dynamics, power structures, communication styles, and social phenomena in online communities, political debates, or interviews.\n",
170+
" * **Customer Service Analytics:** Identifying common customer issues, evaluating agent performance, detecting customer frustration, and optimizing chatbot interactions.\n",
171+
" * **Product Analytics:** Understanding how users interact with a product through chat, forums, or reviews; identifying pain points or feature requests.\n",
172+
" * **Content Moderation:** Detecting toxic language, harassment, or hate speech by analyzing conversational context and patterns, not just isolated words.\n",
173+
" * **Chatbot Development:** Analyzing user input to improve chatbot understanding, evaluating dialog flow, and benchmarking chatbot responses against human conversations.\n",
174+
" * **Digital Humanities:** Analyzing historical texts, interviews, or literary dialogues to uncover linguistic and social insights.\n",
175+
" * **Linguistic Research:** Investigating specific linguistic phenomena, such as politeness strategies, code-switching, or the use of discourse markers in natural conversation.\n",
176+
"\n",
177+
"In essence, ConvoKit empowers researchers and data scientists to move beyond simple text analysis and delve into the rich, structured world of conversations, enabling deeper insights into human communication.\n"
178+
]
179+
}
180+
],
181+
"source": [
182+
"client = get_llm_client(model=MODEL, provider=MODEL_PROVIDER, config_manager=config)\n",
183+
"response = client.generate(\"What can you do with ConvoKit.\")\n",
184+
"print(\"Text:\", response.text)"
185+
]
186+
},
99187
{
100188
"cell_type": "markdown",
101189
"id": "33b72969",
@@ -114,13 +202,13 @@
114202
"outputs": [],
115203
"source": [
116204
"MODEL_PROVIDER = \"gemini\"\n",
117-
"MODEL = \"gemini-2.0-flash-001\"\n",
205+
"MODEL = \"gemini-2.5-flash\"\n",
118206
"config.set_google_cloud_config(\"YOUR PROJECT\", \"YOUR LOCATION\")"
119207
]
120208
},
121209
{
122210
"cell_type": "code",
123-
"execution_count": 28,
211+
"execution_count": null,
124212
"id": "32fad57a",
125213
"metadata": {},
126214
"outputs": [
@@ -136,7 +224,7 @@
136224
}
137225
],
138226
"source": [
139-
"client = get_llm_client(MODEL_PROVIDER, config)\n",
227+
"client = get_llm_client(model=MODEL, provider=MODEL_PROVIDER, config_manager=config, use_vertex_ai=True)\n",
140228
"response = client.generate(\"Tell me a fun fact about Cornell University.\")\n",
141229
"print(\"Text:\", response.text)"
142230
]
@@ -310,7 +398,7 @@
310398
],
311399
"metadata": {
312400
"kernelspec": {
313-
"display_name": "Python 3",
401+
"display_name": "convokit",
314402
"language": "python",
315403
"name": "python3"
316404
},
@@ -324,7 +412,7 @@
324412
"name": "python",
325413
"nbconvert_exporter": "python",
326414
"pygments_lexer": "ipython3",
327-
"version": "3.11.5"
415+
"version": "3.12.11"
328416
}
329417
},
330418
"nbformat": 4,

‎convokit/genai/gemini_client.py‎

Lines changed: 29 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import os
21
from google import genai
32
from google.genai.types import GenerateContentConfig, HttpOptions
43
from .base import LLMClient, LLMResponse
@@ -7,35 +6,48 @@
76

87

98
class GeminiClient(LLMClient):
10-
"""Client for interacting with Google Gemini models via Vertex AI.
9+
"""Client for interacting with Google Gemini models.
1110
12-
This client is configured to use Vertex AI and requires Google Cloud project and location
13-
to be set. Configuration can be provided via the GenAI config system or environment variables.
11+
This client uses the Gemini Developer API with an API key by default. If no API key
12+
is configured, or if ``use_vertex_ai`` is set, it falls back to Vertex AI with Google
13+
Cloud project and location configuration.
1414
15-
:param model: Name of the Gemini model to use (default: "gemini-2.0-flash-001")
15+
:param model: Name of the Gemini model to use
1616
:param config_manager: GenAIConfigManager instance (optional, will create one if not provided)
17+
:param use_vertex_ai: Force Vertex AI authentication instead of API key authentication
1718
"""
1819

1920
def __init__(
2021
self,
21-
model: str = "gemini-2.0-flash-001",
22+
model: str,
2223
config_manager: GenAIConfigManager = None,
24+
use_vertex_ai: bool = False,
2325
):
2426
if config_manager is None:
2527
config_manager = GenAIConfigManager()
2628

2729
self.config_manager = config_manager
30+
self.model = model
31+
32+
api_key = config_manager.get_api_key("gemini")
33+
if api_key and not use_vertex_ai:
34+
self.client = genai.Client(api_key=api_key)
35+
else:
36+
self.client = self._create_vertex_ai_client(config_manager)
2837

29-
# Get required Vertex AI configuration
38+
@staticmethod
39+
def _create_vertex_ai_client(config_manager: GenAIConfigManager):
40+
"""Create a Gemini client authenticated through Vertex AI."""
3041
google_cloud_project = config_manager.get_google_cloud_project()
3142
google_cloud_location = config_manager.get_google_cloud_location()
3243

33-
# Validate required fields
3444
if not google_cloud_project:
3545
raise ValueError(
36-
"Google Cloud project is required for Vertex AI. "
37-
"Set it using config_manager.set_google_cloud_config(project, location) "
38-
"or via GOOGLE_CLOUD_PROJECT environment variable."
46+
"Gemini API key or Google Cloud project is required. "
47+
"Set an API key using config_manager.set_api_key('gemini', 'your-key'), "
48+
"or via GEMINI_API_KEY. For Vertex AI, set project and location using "
49+
"config_manager.set_google_cloud_config(project, location) or via "
50+
"GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION."
3951
)
4052

4153
if not google_cloud_location:
@@ -45,13 +57,12 @@ def __init__(
4557
"or via GOOGLE_CLOUD_LOCATION environment variable."
4658
)
4759

48-
# Set up Vertex AI environment
49-
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "true"
50-
os.environ["GOOGLE_CLOUD_PROJECT"] = google_cloud_project
51-
os.environ["GOOGLE_CLOUD_LOCATION"] = google_cloud_location
52-
53-
self.client = genai.Client(http_options=HttpOptions(api_version="v1"))
54-
self.model = model
60+
return genai.Client(
61+
vertexai=True,
62+
project=google_cloud_project,
63+
location=google_cloud_location,
64+
http_options=HttpOptions(api_version="v1"),
65+
)
5566

5667
def generate(self, prompt, temperature=0.0, times_retried=0) -> LLMResponse:
5768
"""Generate text using the Gemini model.

‎convokit/genai/gpt_client.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,11 @@ class GPTClient(LLMClient):
1111
Handles authentication, request formatting, and error retry logic. API key is managed
1212
through the GenAI config system.
1313
14-
:param model: Name of the GPT model to use (default: "gpt-4o-mini")
14+
:param model: Name of the GPT model to use
1515
:param config_manager: GenAIConfigManager instance (optional, will create one if not provided)
1616
"""
1717

18-
def __init__(self, model: str = "gpt-4o-mini", config_manager: GenAIConfigManager = None):
18+
def __init__(self, model: str, config_manager: GenAIConfigManager = None):
1919
if config_manager is None:
2020
config_manager = GenAIConfigManager()
2121

‎docs/source/genai.rst‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ Supported Providers
4040
Currently supported LLM providers:
4141

4242
* **OpenAI GPT**: Access to OpenAI GPT models through the OpenAI API. See `OpenAI API setup <https://platform.openai.com/docs/quickstart>`_.
43-
* **Google Gemini**: Access to Google Gemini models via Vertex AI. See `Vertex AI setup guide <https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_.
43+
* **Google Gemini**: Access to Google Gemini models via the Gemini Developer API by default, or via Vertex AI. See `Gemini API key setup <https://ai.google.dev/gemini-api/docs/api-key>`_ and `Vertex AI setup guide <https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_.
4444
* **Local Models**: Template implementation for local LLM models (requires custom implementation)
4545

4646
GPT Client
@@ -79,7 +79,7 @@ Configuration
7979
The GenAIConfigManager handles API key storage and retrieval for different LLM providers. It supports:
8080

8181
* **File-based storage**: Configuration is stored in `~/.convokit/config.yml`
82-
* **Environment variables**: API keys can be set via environment variables (e.g., `GPT_API_KEY`)
82+
* **Environment variables**: API keys can be set via environment variables (e.g., `GPT_API_KEY` or `GEMINI_API_KEY`)
8383
* **Secure storage**: API keys are stored locally and not exposed in code
8484
* **Provider-specific settings**: Support for different configuration requirements per provider (e.g., Google Cloud project settings for Gemini)
8585

@@ -94,7 +94,10 @@ The GenAIConfigManager handles API key storage and retrieval for different LLM p
9494
# Set OpenAI API key
9595
config.set_api_key("gpt", "your-openai-api-key")
9696
97-
# Set Google Cloud configuration for Gemini
97+
# Set Gemini API key
98+
config.set_api_key("gemini", "your-gemini-api-key")
99+
100+
# Or set Google Cloud configuration for Gemini on Vertex AI
98101
config.set_google_cloud_config("your-project-id", "your-location")
99102
100103
# Configuration is automatically saved and can be reused

0 commit comments

Comments
 (0)