Skip to content

Commit c8529d0

Browse files
committed
[Docs Agent] Release of Docs Agent v 0.3.2.
What's change: - Support the new Gemini 1.5 models in preview. - Add new experimental CLI features to interact with Gemini models directly from a Linux terminal: `agent tellme` and `agent helpme`. - Better handle uploading of text chunks using the Semantic Retrieval API. - Add a new chat UI feature to provide a page for viewing logs. - Update the Google Generative AI SDK version to `0.5.0`. - Refactor the pre-preprocessing module (in progress). - Remove unused code and type mismatch errors. - Bug fixes.
1 parent 1294973 commit c8529d0

File tree

30 files changed

+3243
-1458
lines changed

30 files changed

+3243
-1458
lines changed

examples/gemini/python/docs-agent/README.md

Lines changed: 40 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,22 @@ The following list summarizes the tasks and features supported by Docs Agent:
7070
- **Run the Docs Agent CLI from anywhere in a terminal**: You can set up the
7171
Docs Agent CLI to ask questions to the Gemini model from anywhere in a terminal.
7272
For more information, see the [Set up Docs Agent CLI][cli-readme] page.
73+
- **Support the Gemini 1.5 models**: You can use the new Gemini 1.5 models,
74+
`gemini-1.5-pro-latest` and `text-embedding-004`, with Docs Agent today.
75+
For the moment, the following `config.yaml` setup is recommended:
76+
77+
```
78+
models:
79+
- language_model: "models/aqa"
80+
embedding_model: "models/text-embedding-004"
81+
api_endpoint: "generativelanguage.googleapis.com"
82+
...
83+
app_mode: "1.5"
84+
db_type: "chroma"
85+
```
86+
87+
The setup above uses 3 Gemini models to their strength: AQA (`aqa`),
88+
Gemini 1.0 Pro (`gemini-pro`), and Gemini 1.5 Pro (`gemini-1.5-pro-latest`).
7389

7490
For more information on Docs Agent's architecture and features,
7591
see the [Docs Agent concepts][docs-agent-concepts] page.
@@ -113,27 +129,24 @@ Update your host machine's environment to prepare for the Docs Agent setup:
113129
2. Install the following dependencies:
114130

115131
```posix-terminal
116-
sudo apt install git pip python3-venv
132+
sudo apt install git pipx python3-venv
117133
```
118134

119135
3. Install `poetry`:
120136

121137
```posix-terminal
122-
curl -sSL https://install.python-poetry.org | python3 -
138+
pipx install poetry
123139
```
124140

125-
**Important**: Make sure that `$HOME/.local/bin` is in your `PATH` variable
126-
(for example, `export PATH=$PATH:~/.local/bin`).
127-
128-
4. Set the following environment variable:
141+
4. To add `$HOME/.local/bin` to your `PATH` variable, run the following
142+
command:
129143

130144
```posix-terminal
131-
export PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring
145+
pipx ensurepath
132146
```
133147

134-
This is a [known issue][poetry-known-issue] in `poetry`.
135-
136-
5. Set the Google API key as a environment variable:
148+
5. To set the Google API key as a environment variable, add the following
149+
line to your `$HOME/.bashrc` file:
137150

138151
```
139152
export GOOGLE_API_KEY=<YOUR_API_KEY_HERE>
@@ -142,8 +155,11 @@ Update your host machine's environment to prepare for the Docs Agent setup:
142155
Replace `<YOUR_API_KEY_HERE>` with the API key to the
143156
[Gemini API][genai-doc-site].
144157

145-
**Tip**: To avoid repeating these `export` lines, add them to your
146-
`$HOME/.bashrc` file.
158+
6. Update your environment:
159+
160+
```posix-termainl
161+
source ~/.bashrc
162+
```
147163

148164
### 3. (Optional) Authorize credentials for Docs Agent
149165

@@ -256,7 +272,7 @@ Update settings in the Docs Agent project to use your custom dataset:
256272

257273
```
258274
inputs:
259-
- path: "/usr/local/home/user01/website/src"
275+
- path: "/usr/local/home/user01/website/src/content"
260276
url_prefix: "https://docs.flutter.dev"
261277
```
262278

@@ -265,30 +281,31 @@ Update settings in the Docs Agent project to use your custom dataset:
265281

266282
```
267283
inputs:
268-
- path: "/usr/local/home/user01/website/src/ui"
284+
- path: "/usr/local/home/user01/website/src/content/ui"
269285
url_prefix: "https://docs.flutter.dev/ui"
270-
- path: "/usr/local/home/user01/website/src/tools"
286+
- path: "/usr/local/home/user01/website/src/content/tools"
271287
url_prefix: "https://docs.flutter.dev/tools"
272288
```
273289

274-
6. (**Optional**) If you want to use the Gemini AQA model and populate a corpus online
275-
via the [Semantic Retrieval API][semantic-api], use the following settings:
290+
6. If you want to use the `gemini-pro` model with a local vector database setup
291+
(`chroma`), use the following settings:
276292

277293
```
278294
models:
279-
- language_model: "models/aqa"
295+
- language_model: "models/gemini-pro"
280296
...
281-
db_type: "google_semantic_retriever"
297+
db_type: "chroma"
282298
```
283299

284-
Or if you want to use the `gemini-pro` model with a local vector database setup
285-
(`chroma`), use the following settings:
300+
(**Optional**) Or if you want to use the Gemini AQA model and populate
301+
a corpus online via the [Semantic Retrieval API][semantic-api], use the
302+
following settings:
286303

287304
```
288305
models:
289-
- language_model: "models/gemini-pro"
306+
- language_model: "models/aqa"
290307
...
291-
db_type: "chroma"
308+
db_type: "google_semantic_retriever"
292309
```
293310

294311
7. Save the `config.yaml` file and exit the text editor.
@@ -403,7 +420,6 @@ Meggin Kearney (`@Meggin`), and Kyo Lee (`@kyolee415`).
403420
[chroma-docs]: https://docs.trychroma.com/
404421
[flutter-docs-src]: https://github.com/flutter/website/tree/main/src
405422
[flutter-docs-site]: https://docs.flutter.dev/
406-
[poetry-known-issue]: https://github.com/python-poetry/poetry/issues/1917
407423
[apps-script-readme]: ./apps_script/README.md
408424
[scripts-readme]: ./docs_agent/preprocess/README.md
409425
[config-yaml]: config.yaml

examples/gemini/python/docs-agent/apps_script/drive_to_markdown.gs

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,11 +44,19 @@ function convertDriveFolderToMDForDocsAgent(folderName) {
4444

4545
while (myfiles.hasNext()) {
4646
var myfile = myfiles.next();
47+
var ftype = myfile.getMimeType();
48+
// If this is a shorcut, retrieve the target file
49+
if (ftype == "application/vnd.google-apps.shortcut") {
50+
var fid = myfile.getTargetId();
51+
var myfile = DriveApp.getFileById(fid);
52+
var ftype = myfile.getMimeType();
53+
}
54+
else{
55+
var fid = myfile.getId();
56+
}
4757
var fname = sanitizeFileName(myfile.getName());
4858
var fdate = myfile.getLastUpdated();
4959
var furl = myfile.getUrl();
50-
var fid = myfile.getId();
51-
var ftype = myfile.getMimeType();
5260
var fcreate = myfile.getDateCreated();
5361

5462
//Function returns an array, assign each array value to seperate variables
@@ -58,7 +66,6 @@ function convertDriveFolderToMDForDocsAgent(folderName) {
5866
var md5_backup = backup_results[1];
5967
var mdoutput_backup_id = backup_results[2];
6068
}
61-
6269
if (ftype == "application/vnd.google-apps.document") {
6370
Logger.log("File: " + fname + " is a Google doc.");
6471
let gdoc = DocumentApp.openById(fid);

examples/gemini/python/docs-agent/apps_script/exportmd.gs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -866,6 +866,8 @@ function processParagraph(index, element, inSrc, imageCounter, listCounters, ima
866866
} else if (t === DocumentApp.ElementType.FOOTNOTE) {
867867
textElements.push(' ('+element.getChild(i).getFootnoteContents().getText()+')');
868868
// Fixes for new elements
869+
} else if (t === DocumentApp.ElementType.EQUATION) {
870+
textElements.push(element.getChild(i).getText());
869871
} else if (t === DocumentApp.ElementType.DATE) {
870872
textElements.push(' ('+element.getChild(i)+')');
871873
} else if (t === DocumentApp.ElementType.RICH_LINK) {
@@ -875,8 +877,8 @@ function processParagraph(index, element, inSrc, imageCounter, listCounters, ima
875877
} else if (t === DocumentApp.ElementType.UNSUPPORTED) {
876878
textElements.push(' <UNSUPPORTED> ');
877879
} else {
878-
throw "Paragraph "+index+" of type "+element.getType()+" has an unsupported child: "
879-
+t+" "+(element.getChild(i)["getText"] ? element.getChild(i).getText():'')+" index="+index;
880+
Logger.log("Paragraph "+index+" of type "+element.getType()+" has an unsupported child: "
881+
+t+" "+(element.getChild(i)["getText"] ? element.getChild(i).getText():'')+" index="+index);
880882
}
881883
}
882884

examples/gemini/python/docs-agent/docs_agent/agents/docs_agent.py

Lines changed: 56 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,26 +16,21 @@
1616

1717
"""Docs Agent"""
1818

19-
import os
20-
import sys
19+
import typing
2120

2221
from absl import logging
2322
import google.api_core
2423
import google.ai.generativelanguage as glm
2524
from chromadb.utils import embedding_functions
2625

27-
from docs_agent.storage.chroma import Chroma, Format, ChromaEnhanced
28-
from docs_agent.models.palm import PaLM
26+
from docs_agent.storage.chroma import ChromaEnhanced
2927

3028
from docs_agent.models.google_genai import Gemini
3129

32-
from docs_agent.utilities.config import ProductConfig, ReadConfig, Input, Models
33-
from docs_agent.models import tokenCount
30+
from docs_agent.utilities.config import ProductConfig, Models
3431
from docs_agent.preprocess.splitters import markdown_splitter
3532

3633
from docs_agent.preprocess.splitters.markdown_splitter import Section as Section
37-
from docs_agent.utilities.helpers import get_project_path
38-
from docs_agent.postprocess.docs_retriever import FullPage as FullPage
3934
from docs_agent.postprocess.docs_retriever import SectionDistance as SectionDistance
4035
from docs_agent.postprocess.docs_retriever import (
4136
SectionProbability as SectionProbability,
@@ -49,6 +44,8 @@ class DocsAgent:
4944
def __init__(self, config: ProductConfig, init_chroma: bool = True):
5045
# Models settings
5146
self.config = config
47+
self.embedding_model = str(self.config.models.embedding_model)
48+
self.api_endpoint = str(self.config.models.api_endpoint)
5249
# Use the new chroma db for all queries
5350
# Should make a function for this or clean this behavior
5451
if init_chroma:
@@ -62,9 +59,9 @@ def __init__(self, config: ProductConfig, init_chroma: bool = True):
6259
)
6360
self.collection = self.chroma.get_collection(
6461
self.collection_name,
65-
embedding_model=self.config.models.embedding_model,
62+
embedding_model=self.embedding_model,
6663
embedding_function=embedding_function_gemini_retrieval(
67-
self.config.models.api_key
64+
self.config.models.api_key, self.embedding_model
6865
),
6966
)
7067
# AQA model settings
@@ -77,9 +74,12 @@ def __init__(self, config: ProductConfig, init_chroma: bool = True):
7774
self.context_model = "models/gemini-pro"
7875
gemini_model_config = Models(
7976
language_model=self.context_model,
80-
embedding_model="models/embedding-001",
77+
embedding_model=self.embedding_model,
78+
api_endpoint=self.api_endpoint,
79+
)
80+
self.gemini = Gemini(
81+
models_config=gemini_model_config, conditions=config.conditions
8182
)
82-
self.gemini = Gemini(models_config=gemini_model_config)
8383
# Semantic retriever
8484
if self.config.db_type == "google_semantic_retriever":
8585
for item in self.config.db_configs:
@@ -93,9 +93,34 @@ def __init__(self, config: ProductConfig, init_chroma: bool = True):
9393
)
9494
self.aqa_response_buffer = ""
9595

96-
if self.config.models.language_model == "models/gemini-pro":
97-
self.gemini = Gemini(models_config=config.models)
98-
self.context_model = "models/gemini-pro"
96+
if self.config.models.language_model.startswith("models/gemini"):
97+
self.gemini = Gemini(
98+
models_config=config.models, conditions=config.conditions
99+
)
100+
self.context_model = self.config.models.language_model
101+
102+
# Always initialize the gemini-pro model for other tasks.
103+
gemini_pro_model_config = Models(
104+
language_model="models/gemini-pro",
105+
embedding_model=self.embedding_model,
106+
api_endpoint=self.api_endpoint,
107+
)
108+
self.gemini_pro = Gemini(
109+
models_config=gemini_pro_model_config, conditions=config.conditions
110+
)
111+
112+
if self.config.app_mode == "1.5":
113+
# Initialize the gemini-1.5.pro model for summarization.
114+
gemini_15_model_config = Models(
115+
language_model="models/gemini-1.5-pro-latest",
116+
embedding_model=self.embedding_model,
117+
api_endpoint=self.api_endpoint,
118+
)
119+
self.gemini_15 = Gemini(
120+
models_config=gemini_15_model_config, conditions=config.conditions
121+
)
122+
else:
123+
self.gemini_15 = self.gemini_pro
99124

100125
# Use this method for talking to a Gemini content model
101126
def ask_content_model_with_context(self, context, question):
@@ -261,7 +286,7 @@ def ask_aqa_model_using_corpora(self, question, answer_style: str = "VERBOSE"):
261286

262287
def ask_aqa_model(self, question):
263288
response = ""
264-
if self.db_type == "ONLINE_STORAGE":
289+
if self.config.db_type == "google_semantic_retriever":
265290
response = self.ask_aqa_model_using_corpora(question)
266291
else:
267292
response = self.ask_aqa_model_using_local_vector_store(question)
@@ -436,7 +461,11 @@ def query_vector_store_to_build(
436461
# If prompt is "fact_checker" it will use the fact_check_question from
437462
# config.yaml for the prompt
438463
def ask_content_model_with_context_prompt(
439-
self, context: str, question: str, prompt: str = None
464+
self,
465+
context: str,
466+
question: str,
467+
prompt: typing.Optional[str] = None,
468+
model: typing.Optional[str] = None,
440469
):
441470
if prompt == None:
442471
prompt = self.config.conditions.condition_text
@@ -447,7 +476,13 @@ def ask_content_model_with_context_prompt(
447476
if self.config.log_level == "VERBOSE":
448477
self.print_the_prompt(new_prompt)
449478
try:
450-
response = self.gemini.generate_content(contents=new_prompt)
479+
response = ""
480+
if model == "gemini-pro":
481+
response = self.gemini_pro.generate_content(contents=new_prompt)
482+
elif model == "gemini-1.5-pro":
483+
response = self.gemini_15.generate_content(contents=new_prompt)
484+
else:
485+
response = self.gemini.generate_content(contents=new_prompt)
451486
except:
452487
return self.config.conditions.model_error_message, new_prompt
453488
for chunk in response:
@@ -475,7 +510,7 @@ def ask_content_model_to_use_file(self, prompt: str, file: str):
475510
# Use this method for asking a Gemini content model for fact-checking.
476511
# This uses ask_content_model_with_context_prompt w
477512
def ask_content_model_to_fact_check_prompt(self, context: str, prev_response: str):
478-
question = self.fact_check_question + "\n\nText: "
513+
question = self.config.conditions.fact_check_question + "\n\nText: "
479514
question += prev_response
480515
return self.ask_content_model_with_context_prompt(
481516
context=context, question=question, prompt=""
@@ -487,7 +522,7 @@ def generate_embedding(self, text, task_type: str = "SEMANTIC_SIMILARITY"):
487522

488523

489524
# Function to give an embedding function for gemini using an API key
490-
def embedding_function_gemini_retrieval(api_key):
525+
def embedding_function_gemini_retrieval(api_key, embedding_model: str):
491526
return embedding_functions.GoogleGenerativeAiEmbeddingFunction(
492-
api_key=api_key, model_name="models/embedding-001", task_type="RETRIEVAL_QUERY"
527+
api_key=api_key, model_name=embedding_model, task_type="RETRIEVAL_QUERY"
493528
)

0 commit comments

Comments
 (0)