Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added create time and update time fields to Corpus, Document, and Chu… #187

Merged
merged 11 commits into from
Jan 31, 2024
Merged

Added create time and update time fields to Corpus, Document, and Chu… #187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
28165d8
Added create time and update time fields to Corpus, Document, and Chu…
shilpakancharla Jan 31, 2024
9b8035a
Fixed test cases
shilpakancharla Jan 31, 2024
640d3cd
Fixed formatting
shilpakancharla Jan 31, 2024
36ea99d
Updated async tests with times
shilpakancharla Jan 31, 2024
80b8854
Formatting with updated black .
shilpakancharla Jan 31, 2024
6e7e822
Merged with main
shilpakancharla Jan 31, 2024
d5b35dc
Testing out changes
shilpakancharla Jan 31, 2024
e262938
Updating type for time objects
shilpakancharla Jan 31, 2024
315e6fb
Resolving Python 3.10 type issue
shilpakancharla Jan 31, 2024
210a94d
Merge branch 'main' into creation_update_time
shilpakancharla Jan 31, 2024
961bcba
Remove _fix_microseconds.
MarkDaoust Jan 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions google/generativeai/types/retriever_types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# limitations under the License.
from __future__ import annotations

import datetime
import re
import string
import abc
Expand Down Expand Up @@ -168,6 +169,8 @@ class Corpus:

name: str
display_name: str
create_time: datetime.datetime
update_time: datetime.datetime

def create_document(
self,
Expand Down Expand Up @@ -498,6 +501,8 @@ class Document(abc.ABC):
name: str
display_name: str
custom_metadata: list[CustomMetadata]
create_time: datetime.datetime
update_time: datetime.datetime

def create_chunk(
self,
Expand Down Expand Up @@ -1138,13 +1143,17 @@ class Chunk(abc.ABC):
data: ChunkData
custom_metadata: list[CustomMetadata] | None
state: State
create_time: datetime.datetime
update_time: datetime.datetime

def __init__(
self,
name: str,
data: ChunkData | str,
custom_metadata: list[CustomMetadata] | None,
state: State,
create_time: datetime.datetime,
update_time: datetime.datetime,
):
self.name = name
if isinstance(data, str):
Expand All @@ -1156,6 +1165,8 @@ def __init__(
else:
self.custom_metadata = [CustomMetadata(*cm) for cm in custom_metadata]
self.state = state
self.create_time = create_time
self.update_time = update_time

def _apply_update(self, path, value):
parts = path.split(".")
Expand Down
86 changes: 76 additions & 10 deletions tests/test_retriever.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,30 +44,55 @@ def create_corpus(
request: glm.CreateCorpusRequest,
) -> glm.Corpus:
self.observed_requests.append(request)
return glm.Corpus(name="corpora/demo_corpus", display_name="demo_corpus")
return glm.Corpus(
name="corpora/demo_corpus",
display_name="demo_corpus",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
def get_corpus(
request: glm.GetCorpusRequest,
) -> glm.Corpus:
self.observed_requests.append(request)
return glm.Corpus(name="corpora/demo_corpus", display_name="demo_corpus")
return glm.Corpus(
name="corpora/demo_corpus",
display_name="demo_corpus",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
def update_corpus(
request: glm.UpdateCorpusRequest,
) -> glm.Corpus:
self.observed_requests.append(request)
return glm.Corpus(name="corpora/demo_corpus", display_name="demo_corpus_1")
return glm.Corpus(
name="corpora/demo_corpus",
display_name="demo_corpus_1",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
def list_corpora(
request: glm.ListCorporaRequest,
) -> glm.ListCorporaResponse:
self.observed_requests.append(request)
return [
glm.Corpus(name="corpora/demo_corpus_1", display_name="demo_corpus_1"),
glm.Corpus(name="corpora/demo_corpus_2", display_name="demo_corpus_2"),
glm.Corpus(
name="corpora/demo_corpus_1",
display_name="demo_corpus_1",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
glm.Corpus(
name="corpora/demo_corpus_2",
display_name="demo_corpus_2",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
]

@add_client_method
Expand All @@ -82,6 +107,8 @@ def query_corpus(
chunk=glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is a demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
)
]
Expand All @@ -97,7 +124,10 @@ def create_document(
) -> retriever_service.Document:
self.observed_requests.append(request)
return glm.Document(
name="corpora/demo_corpus/documents/demo_doc", display_name="demo_doc"
name="corpora/demo_corpus/documents/demo_doc",
display_name="demo_doc",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
Expand All @@ -106,7 +136,10 @@ def get_document(
) -> retriever_service.Document:
self.observed_requests.append(request)
return glm.Document(
name="corpora/demo_corpus/documents/demo_doc", display_name="demo_doc"
name="corpora/demo_corpus/documents/demo_doc",
display_name="demo_doc",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
Expand All @@ -115,7 +148,10 @@ def update_document(
) -> glm.Document:
self.observed_requests.append(request)
return glm.Document(
name="corpora/demo_corpus/documents/demo_doc", display_name="demo_doc_1"
name="corpora/demo_corpus/documents/demo_doc",
display_name="demo_doc_1",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
Expand All @@ -125,10 +161,16 @@ def list_documents(
self.observed_requests.append(request)
return [
glm.Document(
name="corpora/demo_corpus/documents/demo_doc_1", display_name="demo_doc_1"
name="corpora/demo_corpus/documents/demo_doc_1",
display_name="demo_doc_1",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
glm.Document(
name="corpora/demo_corpus/documents/demo_doc_2", display_name="demo_doc_2"
name="corpora/demo_corpus/documents/demo_doc_2",
display_name="demo_doc_2",
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
]

Expand All @@ -150,6 +192,8 @@ def query_document(
chunk=glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is a demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
)
]
Expand All @@ -163,6 +207,8 @@ def create_chunk(
return glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is a demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
Expand All @@ -175,10 +221,14 @@ def batch_create_chunks(
glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/dc",
data={"string_value": "This is a demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/dc1",
data={"string_value": "This is another demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
]
)
Expand All @@ -191,6 +241,8 @@ def get_chunk(
return glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is a demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
Expand All @@ -202,10 +254,14 @@ def list_chunks(
glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is a demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk_1",
data={"string_value": "This is another demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
]

Expand All @@ -215,6 +271,8 @@ def update_chunk(request: glm.UpdateChunkRequest) -> glm.Chunk:
return glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is an updated demo chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
)

@add_client_method
Expand All @@ -227,10 +285,14 @@ def batch_update_chunks(
glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk",
data={"string_value": "This is an updated chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
glm.Chunk(
name="corpora/demo_corpus/documents/demo_doc/chunks/demo_chunk_1",
data={"string_value": "This is another updated chunk."},
create_time="2000-01-01T01:01:01.123456Z",
update_time="2000-01-01T01:01:01.123456Z",
),
]
)
Expand Down Expand Up @@ -301,6 +363,8 @@ def test_query_corpus(self):
"data": {"string_value": "This is a demo chunk."},
"custom_metadata": [],
"state": 0,
"create_time": "2000-01-01T01:01:01.123456Z",
"update_time": "2000-01-01T01:01:01.123456Z",
},
}
]
Expand Down Expand Up @@ -382,6 +446,8 @@ def test_query_document(self):
"data": {"string_value": "This is a demo chunk."},
"custom_metadata": [],
"state": 0,
"create_time": "2000-01-01T01:01:01.123456Z",
"update_time": "2000-01-01T01:01:01.123456Z",
},
}
]
Expand Down
Loading
Loading