fix: multiple documents in embed_documents error (#88)

Fixes #86 --------- Signed-off-by: dependabot[bot] <[email protected]> Signed-off-by: harry-cohere <[email protected]> Co-authored-by: harry-cohere <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Minjie <[email protected]>
langchain-ai · Aug 13, 2024 · e57fee4 · e57fee4
1 parent 16da7a7
commit e57fee4
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 29 deletions.
diff --git a/libs/cohere/langchain_cohere/embeddings.py b/libs/cohere/langchain_cohere/embeddings.py
@@ -1,5 +1,5 @@
 import typing
-from typing import Any, Dict, List, Optional, Sequence
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import cohere
 from langchain_core.embeddings import Embeddings
@@ -133,19 +133,22 @@ def embed(
         *,
         input_type: typing.Optional[cohere.EmbedInputType] = None,
     ) -> List[List[float]]:
-        embeddings = self.embed_with_retry(
+        response = self.embed_with_retry(
             model=self.model,
             texts=texts,
             input_type=input_type,
             truncate=self.truncate,
             embedding_types=self.embedding_types,
         )
-        embeddings = embeddings.dict().get("embeddings", [])
-        return [
-            list(map(float, e[0]))
-            for embedding_type in self.embedding_types
-            if (e := embeddings.get(embedding_type, [[]]))
-        ]
+        embeddings = response.dict().get("embeddings", [])
+        embeddings_as_float: List[List[float]] = []
+        for embedding_type in self.embedding_types:
+            e: List[List[Union[int, float]]] = embeddings.get(embedding_type)
+            if not e:
+                continue
+            for i in range(len(e)):
+                embeddings_as_float.append(list(map(float, e[i])))
+        return embeddings_as_float
 
     async def aembed(
         self,

diff --git a/libs/cohere/poetry.lock b/libs/cohere/poetry.lock
diff --git a/libs/cohere/pyproject.toml b/libs/cohere/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langchain-cohere"
-version = "0.2.1"
+version = "0.2.2"
 description = "An integration package connecting Cohere and LangChain"
 authors = []
 readme = "README.md"

diff --git a/...tests/integration_tests/cassettes/test_langchain_cohere_embedding_multiple_documents.yaml b/...tests/integration_tests/cassettes/test_langchain_cohere_embedding_multiple_documents.yaml
@@ -0,0 +1,74 @@
+interactions:
+- request:
+    body: '{"texts": ["foo bar", "bar foo"], "model": "embed-english-light-v3.0",
+      "input_type": "search_document", "embedding_types": ["float"], "truncate": null}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '151'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.com
+      user-agent:
+      - python-httpx/0.27.0
+      x-client-name:
+      - langchain:partner
+      x-fern-language:
+      - Python
+      x-fern-sdk-name:
+      - cohere
+      x-fern-sdk-version:
+      - 5.5.8
+    method: POST
+    uri: https://api.cohere.com/v1/embed
+  response:
+    body:
+      string: '{"id":"05bbe887-8c3f-4b03-abe3-b3122048a107","texts":["foo bar","bar
+        foo"],"embeddings":{"float":[[-0.023651123,0.021362305,0.023330688,-0.008613586,-0.043640137,0.018920898,0.06744385,-0.021499634,0.14501953,0.072387695,0.040802002,0.027496338,-0.043304443,0.054382324,-0.027862549,0.05731201,-0.0035247803,0.04815674,-0.062927246,-0.011497498,0.095581055,-0.018249512,-0.041809082,-0.0016412735,-0.031463623,-0.016738892,-0.008232117,0.012832642,0.025848389,-0.011741638,-0.021347046,-0.062042236,0.0034275055,0.0027980804,-0.015640259,-0.042999268,0.007293701,0.02053833,0.02330017,0.029846191,-0.025894165,-0.012031555,0.060150146,-0.042541504,0.026382446,0.06970215,0.0012559891,-0.051513672,0.054718018,0.03274536,0.044128418,-0.019714355,0.07336426,0.08648682,0.067993164,0.07873535,0.01902771,0.07080078,0.012802124,0.021362305,-0.07122803,0.03112793,-0.09954834,0.06008911,0.10406494,0.049957275,-0.045684814,-0.0063819885,-0.07336426,0.0395813,-0.057647705,-0.0024681091,0.018493652,0.06958008,-0.02305603,0.015975952,0.03540039,0.034210205,-0.027648926,0.030258179,-0.12585449,-0.024490356,-0.011001587,-0.049987793,0.0074310303,-0.052368164,-0.043884277,-0.07623291,-0.010223389,0.018127441,-0.017028809,-0.0016708374,0.08496094,0.002374649,0.0046844482,0.0022678375,0.015106201,-0.041870117,-0.0065994263,0.044952393,-0.014778137,0.05621338,0.021057129,-0.15185547,0.054870605,0.048187256,0.0041770935,0.03994751,0.016021729,0.023208618,-0.027526855,0.028274536,0.12817383,0.011421204,0.0069122314,-0.11010742,-0.01889038,-0.01977539,-0.14672852,-0.0032672882,0.03552246,0.041992188,0.05432129,0.013923645,-0.03250122,-0.021865845,0.013946533,-0.07922363,0.0032463074,-0.030960083,-0.026504517,0.001531601,0.047943115,0.018814087,0.019607544,0.04763794,0.0049438477,-0.0390625,0.05947876,-0.028274536,-0.019638062,-0.04977417,0.053955078,-0.0005168915,-0.035125732,0.022109985,-0.11407471,-0.021240234,-0.10101318,0.034118652,0.01876831,0.079589844,-0.0004925728,0.026977539,0.0048065186,-0.047576904,0.021514893,0.08746338,-0.03161621,0.0129852295,-0.029144287,0.072631836,-0.01751709,0.10571289,0.0102005005,-0.051696777,0.024261475,-0.023208618,0.024337769,0.001572609,-0.002336502,0.0110321045,0.04888916,0.020874023,-0.014625549,0.008216858,-0.10864258,-0.034484863,-0.006652832,-0.044952393,0.0440979,0.034088135,-0.10015869,0.00945282,0.018981934,-0.01977539,-0.053527832,-0.057403564,0.053771973,0.0038433075,0.017730713,-0.0055656433,-0.023605347,0.05239868,0.03579712,-0.06890869,0.066833496,-0.004360199,-0.050323486,-0.0256958,0.002981186,-0.047424316,-0.08251953,0.05960083,0.009185791,0.03894043,-0.0725708,-0.09039307,0.0029411316,0.023452759,0.034576416,0.034484863,-0.06781006,0.028442383,-0.039855957,-0.11767578,0.07055664,-0.07110596,-0.11743164,0.006275177,0.020233154,-0.011436462,-0.13842773,-0.09552002,0.10748291,0.047546387,-0.01525116,-0.059051514,0.040740967,-0.031402588,0.07836914,-0.020614624,0.022125244,-0.009353638,-0.016098022,0.021499634,-0.037719727,0.050109863,0.07336426,0.03366089,-0.07745361,0.0029029846,-0.014701843,0.028213501,-0.010063171,-0.117004395,0.06561279,0.004432678,0.031707764,-0.021499634,0.13696289,-0.0151901245,-0.049224854,-0.041168213,0.02609253,-0.0071105957,0.07720947,0.11450195,0.0116119385,0.06311035,-0.029800415,-0.041381836,-0.0769043,0.08660889,-0.031234741,-0.004386902,-0.10699463,0.13220215,0.08685303,0.012580872,0.024459839,-0.009902954,-0.061157227,0.022140503,-0.08081055,0.012191772,0.023086548,-0.006767273,-0.06616211,0.0049476624,-0.01058197,-0.00504303,-0.072509766,-0.014144897,-0.03765869,-0.04925537,0.07598877,-0.032287598,0.020812988,-0.018432617,-0.06161499,0.022598267,-0.026428223,0.051452637,-0.053100586,-0.018829346,-0.034301758,0.024261475,-0.08154297,0.022994995,0.057739258,-0.04058838,-0.09069824,0.07092285,-0.047058105,0.06390381,0.04815674,0.007663727,0.008842468,-0.028259277,-0.041381836,-0.026519775,-0.053466797,0.038360596,-0.055908203,0.07543945,0.09429932,-0.06762695,-0.011360168,-0.09466553,-0.04849243,0.012123108,-0.011917114,-0.028579712,0.0049705505,0.0008945465,-0.08880615,-0.10626221,-0.017547607,-0.083984375,0.04849243,0.050476074,0.030395508,-0.07824707,0.033966064,0.014022827,0.041046143,-0.050231934,-0.0046653748,-0.010467529,-0.0053520203,0.06286621,-0.034606934,0.020874023,0.089782715,-0.0135269165,-0.032287598,0.05303955,0.0033226013,-0.049102783,0.0038375854,0.10424805,0.04244995,0.02670288,0.07824707,-0.025787354,-0.0463562,0.055358887,-0.050201416,0.13171387,0.026489258,0.026687622,0.05682373,-0.061065674,0.044830322,-0.03289795,0.016616821,0.011177063,0.0019521713,-0.07562256,-0.014503479,0.009414673,0.04574585,-0.00365448,-0.028411865,-0.056030273,-0.010650635,-0.0009794235,0.014709473,-0.039916992,-0.04550171,0.010017395,-0.02507019,-0.0007619858],[-0.02130127,0.025558472,0.025054932,-0.010940552,-0.04437256,0.0039901733,0.07562256,-0.02897644,0.14465332,0.06951904,0.03253174,0.012428284,-0.052886963,0.068603516,-0.033111572,0.05154419,-0.005168915,0.049468994,-0.06427002,-0.00006866455,0.07763672,-0.013015747,-0.048339844,0.0018844604,-0.011245728,-0.028533936,-0.017959595,0.020019531,0.02859497,0.002292633,-0.0052223206,-0.06921387,-0.01675415,0.0059394836,-0.017593384,-0.03363037,0.0006828308,0.0032958984,0.018692017,0.02293396,-0.015930176,-0.0047073364,0.068725586,-0.039978027,0.026489258,0.07501221,-0.014595032,-0.051696777,0.058502197,0.029388428,0.032806396,0.0049324036,0.07910156,0.09692383,0.07598877,0.08203125,0.010353088,0.07086182,0.022201538,0.029724121,-0.0847168,0.034423828,-0.10620117,0.05505371,0.09466553,0.0463562,-0.05706787,-0.0053520203,-0.08288574,0.035583496,-0.043060303,0.008628845,0.0231781,0.06225586,-0.017074585,0.0052337646,0.036102295,0.040527344,-0.03338623,0.031204224,-0.1282959,-0.013534546,-0.022064209,-0.06488037,0.03173828,-0.03829956,-0.036834717,-0.0748291,0.0072135925,0.027648926,-0.03955078,0.0025348663,0.095336914,-0.0036334991,0.021377563,0.011131287,0.027008057,-0.03930664,-0.0077209473,0.044128418,-0.025817871,0.06829834,0.035461426,-0.14404297,0.05319214,0.04916382,0.011024475,0.046173096,0.030731201,0.028244019,-0.035491943,0.04196167,0.1274414,0.0052719116,0.024353027,-0.101623535,-0.014862061,-0.027145386,-0.13061523,-0.009010315,0.049072266,0.041259766,0.039642334,0.022949219,-0.030838013,-0.02142334,0.018753052,-0.079956055,-0.00894928,-0.027648926,-0.020889282,-0.022003174,0.060150146,0.034698486,0.017547607,0.050689697,0.006504059,-0.018707275,0.059814453,-0.047210693,-0.032043457,-0.060638428,0.064453125,-0.012718201,-0.02218628,0.010734558,-0.10412598,-0.021148682,-0.097229004,0.053894043,0.0044822693,0.06506348,0.0027389526,0.022323608,0.012184143,-0.04815674,0.01828003,0.09777832,-0.016433716,0.011360168,-0.031799316,0.056121826,-0.0135650635,0.10449219,0.0046310425,-0.040740967,0.033996582,-0.04940796,0.031585693,0.008346558,0.0077934265,0.014282227,0.02444458,0.025115967,-0.010910034,0.0027503967,-0.109069824,-0.047424316,-0.0023670197,-0.048736572,0.05947876,0.037231445,-0.11230469,0.017425537,-0.0032978058,-0.006061554,-0.06542969,-0.060028076,0.032714844,-0.0023231506,0.01966858,0.01878357,-0.02243042,0.047088623,0.037475586,-0.05960083,0.04711914,-0.025405884,-0.04724121,-0.013458252,0.014450073,-0.053100586,-0.10595703,0.05899048,0.013031006,0.027618408,-0.05999756,-0.08996582,0.004890442,0.03845215,0.04815674,0.033569336,-0.06359863,0.022140503,-0.025558472,-0.12524414,0.07763672,-0.07550049,-0.09112549,0.0050735474,0.011558533,-0.012008667,-0.11541748,-0.09399414,0.11071777,0.042053223,-0.024887085,-0.058135986,0.044525146,-0.027145386,0.08154297,-0.017959595,0.01826477,-0.0044555664,-0.022949219,0.03326416,-0.04827881,0.051116943,0.082214355,0.031463623,-0.07745361,-0.014221191,-0.009307861,0.03314209,0.004562378,-0.11480713,0.09362793,0.0027122498,0.04586792,-0.02444458,0.13989258,-0.024658203,-0.044036865,-0.041931152,0.025009155,-0.011001587,0.059936523,0.09039307,0.016586304,0.074279785,-0.012687683,-0.034332275,-0.077819824,0.07891846,-0.029006958,-0.0038642883,-0.11590576,0.12432861,0.101623535,0.027709961,0.023086548,-0.013420105,-0.068847656,0.011184692,-0.05999756,0.0041656494,0.035095215,-0.013175964,-0.06488037,-0.011558533,-0.018371582,0.012779236,-0.085632324,-0.02696228,-0.064453125,-0.05618286,0.067993164,-0.025558472,0.027542114,-0.015235901,-0.076416016,0.017700195,-0.0059547424,0.038635254,-0.062072754,-0.025115967,-0.040863037,0.0385437,-0.06500244,0.015731812,0.05581665,-0.0574646,-0.09466553,0.06341553,-0.044769287,0.08337402,0.045776367,0.0151901245,0.008422852,-0.048339844,-0.0368042,-0.015991211,-0.063964844,0.033447266,-0.050476074,0.06463623,0.07019043,-0.06573486,-0.0129852295,-0.08319092,-0.05038452,0.0048713684,-0.0034999847,-0.03050232,-0.003364563,-0.0036258698,-0.064453125,-0.09649658,-0.01852417,-0.08691406,0.043182373,0.03945923,0.033691406,-0.07531738,0.033111572,0.0048065186,0.023880005,-0.05206299,-0.014328003,-0.015899658,0.010322571,0.050720215,-0.028533936,0.023757935,0.089416504,-0.020568848,-0.020843506,0.037231445,0.00022995472,-0.04458618,-0.023025513,0.10076904,0.047180176,0.035614014,0.072143555,-0.020324707,-0.02230835,0.05899048,-0.055419922,0.13513184,0.035369873,0.008255005,0.039855957,-0.068847656,0.031555176,-0.025253296,0.00623703,0.0059890747,0.017578125,-0.07495117,-0.0079956055,0.008033752,0.06530762,-0.027008057,-0.04309082,-0.05218506,-0.016937256,-0.009376526,0.004360199,-0.04888916,-0.039367676,0.0021629333,-0.019958496,-0.0036334991]]},"meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":4}},"response_type":"embeddings_by_type"}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Transfer-Encoding:
+      - chunked
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-type:
+      - application/json
+      date:
+      - Mon, 12 Aug 2024 15:58:52 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      num_chars:
+      - '14'
+      num_tokens:
+      - '4'
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      vary:
+      - Origin
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - 214d2c3dae6f0d8c9c8863027474cc17
+      x-envoy-upstream-service-time:
+      - '31'
+      x-trial-endpoint-call-limit:
+      - '100'
+      x-trial-endpoint-call-remaining:
+      - '99'
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/libs/cohere/tests/integration_tests/test_embeddings.py b/libs/cohere/tests/integration_tests/test_embeddings.py
@@ -20,6 +20,16 @@ def test_langchain_cohere_embedding_documents() -> None:
     assert len(output[0]) > 0
 
 
+@pytest.mark.vcr()
+def test_langchain_cohere_embedding_multiple_documents() -> None:
+    documents = ["foo bar", "bar foo"]
+    embedding = CohereEmbeddings(model="embed-english-light-v3.0")
+    output = embedding.embed_documents(documents)
+    assert len(output) == 2
+    assert len(output[0]) > 0
+    assert len(output[1]) > 0
+
+
 @pytest.mark.vcr()
 def test_langchain_cohere_embedding_query() -> None:
     document = "foo bar"