From 5e966f5c8beaebf6045e76278b3bd97993cd9892 Mon Sep 17 00:00:00 2001 From: David Kristek Date: Thu, 8 Feb 2024 15:53:48 +0100 Subject: [PATCH 1/2] feat(llama-index): add embeddings Refs: #311 --- .../llama_index/llama_index_embedding.py | 34 ++++++ src/genai/extensions/llama_index/__init__.py | 3 +- .../extensions/llama_index/embeddings.py | 53 +++++++++ ...gChainEmbeddings.test_async_embedding.yaml | 111 ++++++++++++++++++ ...estLangChainEmbeddings.test_embedding.yaml | 111 ++++++++++++++++++ .../extensions/test_llama_index_embeddings.py | 39 ++++++ 6 files changed, 350 insertions(+), 1 deletion(-) create mode 100644 examples/extensions/llama_index/llama_index_embedding.py create mode 100644 src/genai/extensions/llama_index/embeddings.py create mode 100644 tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml create mode 100644 tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml create mode 100644 tests/integration/extensions/test_llama_index_embeddings.py diff --git a/examples/extensions/llama_index/llama_index_embedding.py b/examples/extensions/llama_index/llama_index_embedding.py new file mode 100644 index 00000000..d33082eb --- /dev/null +++ b/examples/extensions/llama_index/llama_index_embedding.py @@ -0,0 +1,34 @@ +""" +Llamaindex Embeddings +""" +from dotenv import load_dotenv + +from genai import Client, Credentials +from genai.extensions.llama_index import LlamaIndexEmbeddingsInterface +from genai.schema import TextEmbeddingParameters + +# make sure you have a .env file under genai root with +# GENAI_KEY= +# GENAI_API= (optional) DEFAULT_API = "https://bam-api.res.ibm.com" +load_dotenv() + + +def heading(text: str) -> str: + """Helper function for centering text.""" + return "\n" + f" {text} ".center(80, "=") + "\n" + + +print(heading("LlamaIndex Embeddings")) + +client = Client(credentials=Credentials.from_env()) +embeddings = LlamaIndexEmbeddingsInterface( + client=client, + model_id="sentence-transformers/all-minilm-l6-v2", + parameters=TextEmbeddingParameters(truncate_input_tokens=True), +) + +query_embedding = embeddings.get_query_embedding("Hello world!") +print(query_embedding) + +documents_embedding = embeddings.get_agg_embedding_from_queries(["First document", "Second document"]) +print(documents_embedding) diff --git a/src/genai/extensions/llama_index/__init__.py b/src/genai/extensions/llama_index/__init__.py index 40ffc6a6..80a68d68 100644 --- a/src/genai/extensions/llama_index/__init__.py +++ b/src/genai/extensions/llama_index/__init__.py @@ -1,5 +1,6 @@ """Extension for LLamaIndex library""" +from genai.extensions.llama_index.embeddings import LlamaIndexEmbeddingsInterface from genai.extensions.llama_index.llm import IBMGenAILlamaIndex -__all__ = ["IBMGenAILlamaIndex"] +__all__ = ["IBMGenAILlamaIndex", "LlamaIndexEmbeddingsInterface"] diff --git a/src/genai/extensions/llama_index/embeddings.py b/src/genai/extensions/llama_index/embeddings.py new file mode 100644 index 00000000..2e5bdd6c --- /dev/null +++ b/src/genai/extensions/llama_index/embeddings.py @@ -0,0 +1,53 @@ +import asyncio +from typing import Optional + +from genai._types import ModelLike +from genai.client import Client +from genai.schema import TextEmbeddingParameters +from genai.text.embedding.embedding_service import CreateExecutionOptions + +try: + from llama_index.embeddings.base import BaseEmbedding, Embedding +except ImportError: + raise ImportError("Could not import llamaindex: Please install ibm-generative-ai[llama-index] extension.") # noqa: B904 + + +class LlamaIndexEmbeddingsInterface(BaseEmbedding): + client: Client + model_id: str + parameters: Optional[ModelLike[TextEmbeddingParameters]] = None + execution_options: Optional[ModelLike[CreateExecutionOptions]] = None + + @classmethod + def class_name(cls) -> str: + return "IBMGenAIEmbedding" + + def _get_query_embedding(self, query: str) -> Embedding: + response = self._get_embeddings([query]) + return response[0] + + def _get_text_embedding(self, text: str) -> Embedding: + response = self._get_embeddings([text]) + return response[0] + + def _get_text_embeddings(self, texts: list[str]) -> list[Embedding]: + response = self._get_embeddings(texts) + return response + + async def _aget_query_embedding(self, query: str) -> Embedding: + return await asyncio.get_running_loop().run_in_executor(None, self._get_query_embedding, query) + + async def _aget_text_embedding(self, text: str) -> Embedding: + return await asyncio.get_running_loop().run_in_executor(None, self._get_text_embedding, text) + + async def _aget_text_embeddings(self, texts: list[str]) -> list[Embedding]: + return await asyncio.get_running_loop().run_in_executor(None, self._get_text_embeddings, texts) + + def _get_embeddings(self, texts: list[str]) -> list[Embedding]: + embeddings: list[list[float]] = [] + for response in self.client.text.embedding.create( + model_id=self.model_id, inputs=texts, parameters=self.parameters, execution_options=self.execution_options + ): + embeddings.extend(response.results) + + return embeddings diff --git a/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml new file mode 100644 index 00000000..2b189c1b --- /dev/null +++ b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml @@ -0,0 +1,111 @@ +interactions: +- request: + body: '' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + method: GET + uri: https://api.com/v2/text/embeddings/limits?version=2023-11-22 + response: + body: + string: '{"result":{"concurrency":{"limit":20,"remaining":20}}}' + headers: + cache-control: + - private + content-length: + - '54' + content-type: + - application/json; charset=utf-8 + content-version: + - '2023-11-22' + date: + - Thu, 08 Feb 2024 14:04:40 GMT + keep-alive: + - timeout=72 + set-cookie: + - 2eef5f4c257f6bca76e8da5586743beb=8b795e933fe23741a5fd09cc4f1ebc07; path=/; + HttpOnly; Secure; SameSite=None + vary: + - accept-encoding + status: + code: 200 + message: OK +- request: + body: '{"input": "First document", "model_id": "sentence-transformers/all-minilm-l6-v2", + "parameters": {"truncate_input_tokens": true}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '128' + content-type: + - application/json + method: POST + uri: https://api.com/v2/text/embeddings?version=2023-11-22 + response: + body: + string: '{"results":[[-0.07611478865146637,0.053286295384168625,0.011630389839410782,0.02703363634645939,-0.020964650437235832,-0.0017028871225193143,-0.02678748220205307,0.04056990146636963,0.04763270914554596,0.014788346365094185,0.04339343681931496,0.07593807578086853,0.0012843846343457699,0.00855201669037342,-0.0745568573474884,0.006555262487381697,-0.018848007544875145,-0.0644168108701706,0.03469758480787277,0.06712162494659424,0.0019739633426070213,0.060395244508981705,-0.005633489694446325,-0.027578119188547134,0.017701366916298866,0.028293320909142494,-0.11111179739236832,0.0075835236348211765,0.04227476194500923,-0.11056772619485855,0.0488184429705143,-0.05738593265414238,0.09648681432008743,0.044576071202754974,0.05660592019557953,0.009718634188175201,0.05797331780195236,-0.05440652742981911,0.036171287298202515,0.06543885916471481,-0.0022649781312793493,-0.08455892652273178,-0.029309898614883423,-0.015962744131684303,0.01313595287501812,-0.0004417791496962309,0.010334083810448647,0.0992552787065506,0.049126140773296356,0.02620592713356018,-0.02891329675912857,-0.03497561067342758,-0.02703670598566532,0.018123308196663857,0.057025808840990067,0.04752988740801811,-0.015956465154886246,0.06666743755340576,-0.056488726288080215,-0.06855082511901855,-0.05432615429162979,0.007647170685231686,-0.07357412576675415,0.03254223242402077,0.12377501279115677,0.028229698538780212,-0.004680067300796509,-0.00489498907700181,-0.04303337633609772,0.04150404408574104,-0.016368089243769646,0.0326528362929821,0.01849137246608734,0.030499456450343132,-0.015698838979005814,-0.07027243822813034,0.006100979633629322,0.0350557342171669,0.04501684010028839,-0.09169086068868637,0.004178141709417105,-0.022076696157455444,0.030321242287755013,0.024455692619085312,-0.11486124247312546,-0.0025273847859352827,-0.019262436777353287,0.022405091673135757,-0.005074415821582079,-0.015164084732532501,0.007719150744378567,-0.04771441966295242,0.05346492677927017,0.044002655893564224,-0.07894007861614227,0.045628711581230164,-0.00509940879419446,-0.01611689105629921,0.10104403644800186,0.15719443559646606,0.0031960131600499153,0.023737283423542976,0.1408112645149231,0.06969963759183884,0.006006232928484678,-0.08215378224849701,-0.051280129700899124,-0.04340577498078346,-0.026042301207780838,-0.04198248311877251,-0.008222031407058239,-0.04777222126722336,-0.10308240354061127,0.013990739360451698,-0.04264850914478302,0.002755384426563978,0.031768303364515305,0.05728062614798546,0.02951638586819172,0.0031079030595719814,0.024997133761644363,0.010455358773469925,-0.02089112438261509,-0.030192524194717407,-0.1047578677535057,-0.0630701556801796,0.05935980752110481,-3.689310325890858e-33,0.018732953816652298,-0.00007181623368524015,-0.07653958350419998,0.08046584576368332,-0.03137052059173584,0.016886163502931595,0.01065531000494957,0.026087334379553795,-0.13024289906024933,0.021207323297858238,0.010911231860518456,0.010168812237679958,0.0050546638667583466,0.006856250576674938,-0.08671053498983383,0.034387242048978806,-0.0887618139386177,0.03489241376519203,-0.047422315925359726,-0.05843717232346535,-0.05840500816702843,-0.010714150965213776,0.05513739958405495,-0.004156034905463457,0.01576279290020466,0.033243175595998764,0.011517849750816822,-0.08593405783176422,-0.03567929565906525,-0.0134825948625803,-0.057146165519952774,-0.01318218745291233,-0.026386702433228493,-0.0748041421175003,-0.023853925988078117,0.10301246494054794,0.04259973391890526,-0.044059813022613525,-0.014225777238607407,-0.03309769555926323,-0.028882132843136787,0.05972263216972351,0.054447341710329056,-0.03648264333605766,-0.025904260575771332,0.032167814671993256,-0.05100731551647186,0.07830699533224106,0.12723541259765625,0.018466690555214882,-0.0486471951007843,0.017355622723698616,-0.045035943388938904,0.03343017399311066,-0.0060922387056052685,-0.03044017218053341,-0.0040046460926532745,-0.008548876270651817,-0.008502165786921978,-0.030725140124559402,0.018580880016088486,0.06358906626701355,-0.024671146646142006,0.06750013679265976,-0.013552273623645306,-0.048398926854133606,-0.031233591958880424,0.018136467784643173,0.09101729094982147,-0.08552554994821548,-0.10574617236852646,-0.03554975241422653,0.023665929213166237,-0.09693295508623123,0.0900241956114769,0.0746527761220932,0.07275242358446121,-0.056399762630462646,-0.060314252972602844,-0.04025064781308174,-0.09734656661748886,-0.0005814050091430545,-0.013118426315486431,0.0647367462515831,-0.029480375349521637,0.09067367017269135,0.08380653709173203,-0.0018367903539910913,-0.006479343399405479,-0.007722273003309965,0.00932267401367426,0.03696272149682045,0.05342373996973038,-0.018204154446721077,0.08472762256860733,2.457529632131353e-33,0.006457834504544735,-0.055136796087026596,0.006394235882908106,0.04257359355688095,0.040424179285764694,0.030993953347206116,0.04744979366660118,0.02580159716308117,-0.010715586133301258,0.05614149942994118,0.07997943460941315,-0.012969913892447948,0.04406759887933731,-0.007733637932687998,0.01743471808731556,-0.024840814992785454,0.05109093710780144,-0.032808005809783936,0.03324175998568535,0.051604095846414566,-0.01836460828781128,-0.04881879314780235,-0.03895711898803711,0.06022154167294502,0.018858902156352997,0.01531311683356762,0.07880528271198273,-0.057667337357997894,-0.07734047621488571,-0.029612315818667412,-0.017924057319760323,-0.007808707654476166,-0.005763333290815353,0.07769045233726501,0.013317340053617954,-0.004131403751671314,-0.006829418241977692,-0.006058764178305864,0.020915012806653976,0.08015037328004837,0.06981609016656876,0.022621115669608116,-0.011095408350229263,-0.0010161078535020351,-0.047971002757549286,-0.06444189697504044,0.024511374533176422,0.02654142677783966,0.0009377644746564329,-0.042746491730213165,-0.03214338794350624,0.04735948517918587,0.022207461297512054,-0.14833605289459229,-0.012862633913755417,0.001661280868574977,-0.01834721677005291,-0.03249534219503403,0.0025780177675187588,0.07029899209737778,0.03629522770643234,0.08275070786476135,-0.06574779003858566,0.047293007373809814,-0.042443957179784775,-0.09515932947397232,-0.07330439239740372,0.05922340974211693,-0.0368594266474247,0.01581326127052307,-0.041240546852350235,0.02242102473974228,0.007569017820060253,-0.05807928368449211,0.05175217613577843,-0.017240695655345917,0.033588897436857224,-0.03882821276783943,-0.0032414624001830816,-0.07458201050758362,0.060188066214323044,0.05588485673069954,-0.07107950747013092,0.007625015452504158,0.09588012099266052,-0.0013106901897117496,0.05206756293773651,-0.11739741265773773,-0.032384924590587616,0.028821000829339027,-0.10565381497144699,0.040773868560791016,0.05618344619870186,0.06782420724630356,-0.08374620229005814,-1.2194121268294111e-8,-0.10979562252759933,0.014475689269602299,0.006779859308153391,0.018453504890203476,0.018357109278440475,0.04331545531749725,0.025635426864027977,0.023104725405573845,-0.08535239100456238,-0.10341934859752655,0.01501353457570076,0.03674392029643059,0.04435298219323158,-0.05127701163291931,-0.011615892872214317,-0.01459368597716093,0.0279467161744833,-0.10243868082761765,-0.05912458151578903,0.004412763752043247,0.022360317409038544,0.013892627321183681,0.0198416318744421,-0.0965014323592186,0.028432726860046387,0.003722630674019456,0.04970787465572357,0.019583890214562416,0.03422922641038895,-0.0013914067531004548,-0.027568506076931953,0.10341252386569977,-0.01579933986067772,0.024184221401810646,0.05524905025959015,0.03001238778233528,0.0234824325889349,0.06176608055830002,-0.03289225324988365,-0.06597874313592911,0.021079285070300102,0.026373285800218582,-0.04654461145401001,-0.01519238855689764,-0.032833341509103775,0.003992781043052673,0.0065460167825222015,-0.05169658362865448,0.021103447303175926,-0.10728193819522858,-0.02706659585237503,-0.06279673427343369,0.1034388542175293,0.056007225066423416,-0.02288142964243889,0.03683370351791382,0.06275887787342072,0.013564606197178364,-0.05565091222524643,0.03672057390213013,0.03527916222810745,0.066694475710392,0.09273931384086609,0.027666479349136353]]}' + headers: + content-length: + - '8058' + content-type: + - application/json; charset=utf-8 + content-version: + - '2023-11-22' + date: + - Thu, 08 Feb 2024 14:04:41 GMT + keep-alive: + - timeout=72 + transfer-encoding: + - chunked + vary: + - accept-encoding + status: + code: 200 + message: OK +- request: + body: '{"input": "Second document", "model_id": "sentence-transformers/all-minilm-l6-v2", + "parameters": {"truncate_input_tokens": true}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '129' + content-type: + - application/json + method: POST + uri: https://api.com/v2/text/embeddings?version=2023-11-22 + response: + body: + string: '{"results":[[-0.08104048669338226,0.0376598946750164,0.0020215765107423067,0.026549046859145164,-0.007370871026068926,0.011257929727435112,-0.02854747697710991,0.0631430447101593,0.03664477914571762,-0.03365236148238182,0.06271136552095413,0.06864234060049057,0.02026943676173687,-0.030009351670742035,-0.03508774936199188,0.0023406269028782845,-0.08050798624753952,-0.016443617641925812,0.0018381527625024319,0.062218863517045975,0.03620733320713043,0.009204098023474216,-0.015127652324736118,-0.025030456483364105,0.010739271529018879,0.06100801005959511,-0.11221961677074432,0.05921312794089317,0.0026483291294425726,-0.10150498896837234,0.05173927918076515,0.023861411958932877,-0.033372968435287476,-0.001777892466634512,0.021250296384096146,-0.058816056698560715,0.018938081339001656,-0.00041302599129267037,0.043738409876823425,-0.004523207433521748,-0.011076416820287704,-0.06403172016143799,-0.05241154134273529,-0.04321694001555443,-0.017592426389455795,0.011236284859478474,-0.04660659283399582,0.07122266292572021,0.044390350580215454,0.0497487410902977,-0.07200715690851212,-0.04020734876394272,-0.06696461886167526,0.029645787551999092,0.0457756444811821,0.11751767247915268,-0.012403221800923347,0.056834425777196884,-0.04663803055882454,-0.06756290793418884,0.0081444401293993,0.05049721896648407,-0.10587693005800247,0.0636945515871048,0.1167907789349556,-0.012800364755094051,-0.016207052394747734,-0.004366802517324686,-0.11997950077056885,0.026186034083366394,0.014950335957109928,0.03286281228065491,0.07058639079332352,-0.03315839171409607,0.06088574603199959,-0.10461709648370743,-0.02542501501739025,0.025162100791931152,0.014831612817943096,-0.15255455672740936,0.02012590691447258,0.010713051073253155,0.014526871033012867,-0.04263008385896683,-0.08039771765470505,-0.0731976181268692,-0.00012689009599853307,0.0034983744844794273,-0.022338520735502243,-0.006795673631131649,0.0667412057518959,-0.055777598172426224,0.11562377214431763,0.034403957426548004,-0.06897035986185074,0.0018935615662485361,-0.002328364411368966,0.044492337852716446,0.11536921560764313,0.14253538846969604,-0.017649555578827858,0.00923872459679842,0.08490171283483505,-0.00015613385767210275,0.04251273721456528,-0.09667838364839554,-0.0794786885380745,-0.06583905220031738,0.037441436201334,-0.061086781322956085,0.02558034099638462,0.014064162038266659,-0.09667671471834183,0.026026040315628052,0.01610356941819191,0.06567886471748352,0.08934935927391052,0.05728030577301979,0.03825286403298378,-0.004467225167900324,-0.03369850665330887,-0.011704717762768269,0.015539701096713543,-0.060387320816516876,-0.06724131852388382,-0.10934903472661972,0.07420971244573593,-3.517199218489259e-33,-0.0010445192456245422,-0.04707157239317894,-0.09643448889255524,0.061116773635149,0.013092653825879097,0.048495952039957047,0.012917639687657356,0.05141948536038399,-0.11057691276073456,-0.02722175046801567,-0.013994114473462105,0.014674837701022625,0.007228940725326538,0.0286633912473917,-0.06814133375883102,0.002892223419621587,-0.013423651456832886,0.09643645584583282,0.0409725159406662,-0.012952705845236778,0.009580191224813461,0.06898463517427444,0.0559326708316803,0.013785253278911114,0.05769336596131325,0.07820483297109604,0.04037198796868324,-0.06474727392196655,-0.045225292444229126,-0.002468392252922058,0.031135447323322296,-0.003809040179476142,0.016344018280506134,-0.030952706933021545,0.0029499472584575415,0.07108024507761002,0.020287014544010162,-0.027239948511123657,0.01470908336341381,-0.03253556415438652,-0.02227545529603958,0.023334389552474022,0.04677804186940193,-0.02296368218958378,0.01581859588623047,-0.012426641769707203,-0.009094944223761559,0.030619749799370766,0.12037337571382523,-0.01623496040701866,-0.020219366997480392,0.026702173054218292,-0.056418877094984055,0.020353808999061584,0.03395252302289009,-0.009550661779940128,-0.0622478723526001,0.022805457934737206,0.062185581773519516,0.001912244246341288,0.0438990518450737,0.08762244880199432,-0.025203194469213486,0.04708396643400192,-0.0020530272740870714,-0.003188277827575803,-0.08097474277019501,0.0046669184230268,0.09208934754133224,-0.03296930342912674,-0.10864203423261642,-0.07346982508897781,0.06713204085826874,-0.0824589729309082,0.08118608593940735,-0.015304031781852245,0.03169460967183113,-0.05894801393151283,-0.036768555641174316,-0.052348583936691284,-0.13439710438251495,-0.019004182890057564,-0.029325569048523903,0.0389697328209877,-0.03179283067584038,0.07625920325517654,0.029881611466407776,-0.0287876445800066,-0.03186427801847458,0.051656320691108704,-0.020476073026657104,0.06108597293496132,-0.03309909999370575,-0.03133446350693703,0.14522379636764526,1.9997601731684913e-33,-0.02454514242708683,-0.02168901264667511,-0.04479790851473808,0.04633335769176483,0.033506955951452255,0.039710696786642075,0.04203144833445549,-0.025988228619098663,-0.059948988258838654,0.04111616685986519,0.04170918092131615,-0.010260396637022495,0.08616381138563156,-0.01634076237678528,-0.023933496326208115,-0.03121459297835827,0.053647492080926895,0.009220186620950699,0.013340904377400875,0.0437324121594429,-0.02444729581475258,-0.03904367983341217,0.030613556504249573,0.01505532581359148,0.05224118381738663,-0.002882804023101926,0.011661234311759472,-0.0403885692358017,-0.05561037361621857,-0.03922025114297867,-0.017746806144714355,-0.024655308574438095,-0.04797155410051346,0.03661811724305153,0.046329937875270844,-0.03573417663574219,0.0007827881490811706,-0.032068975269794464,-0.04048440232872963,0.0529610775411129,0.06700104475021362,0.05207323282957077,0.025980258360505104,0.003316115355119109,-0.05455101653933525,0.021097073331475258,0.015427522361278534,0.012054149992763996,0.0036217537708580494,-0.047605108469724655,-0.0024289479479193687,-0.023437868803739548,-0.016562774777412415,-0.09189718216657639,-0.06719052046537399,0.009031710214912891,0.015131138265132904,-0.018840065225958824,0.03931032866239548,0.04334070906043053,0.09708531945943832,0.04747961834073067,-0.06743241101503372,0.03180370107293129,0.017487889155745506,-0.08593153208494186,-0.07724672555923462,0.02401409111917019,-0.03559126332402229,0.07891940325498581,-0.11396396905183792,0.000023013402824290097,0.03012612834572792,-0.05196437984704971,0.05744696781039238,0.00019338278798386455,0.02741345390677452,-0.05575890466570854,-0.008595319464802742,-0.051599353551864624,0.06556955724954605,0.03357896953821182,-0.02616681344807148,0.022066596895456314,-0.0033872537314891815,-0.000342347368132323,0.020659705623984337,-0.03631681948900223,-0.03139133006334305,-0.01155750174075365,-0.06169193983078003,0.08895165473222733,-0.013894675299525261,0.012268207967281342,-0.03348323702812195,-1.2230937151969101e-8,-0.06460992246866226,-0.017255505546927452,0.036914072930812836,-0.036867473274469376,0.005709514953196049,-0.023781590163707733,-0.008763302117586136,0.0318242646753788,-0.06689297407865524,-0.06356614083051682,0.012779262848198414,0.01958351768553257,0.06348159909248352,-0.03766090050339699,-0.034059420228004456,-0.04260662943124771,0.07547925412654877,-0.10275367647409439,-0.04524024948477745,-0.03890029340982437,0.026446901261806488,-0.009833211079239845,0.03849169984459877,-0.009132797829806805,-0.034811895340681076,0.05913154408335686,0.04117123410105705,-0.015952495858073235,-0.013833853416144848,-0.051513154059648514,0.017751075327396393,0.06703679263591766,-0.03192468360066414,0.02006058767437935,0.0793311819434166,-0.004208602476865053,0.045045576989650726,0.07222826033830643,-0.04168285056948662,0.0011332163121551275,0.04194065183401108,-0.024034954607486725,-0.053235992789268494,0.05583583563566208,0.08766256272792816,-0.014262460172176361,0.0411042682826519,-0.10182081907987595,0.029361538589000702,-0.04424484819173813,0.01873917318880558,-0.0940093845129013,0.1568623185157776,0.050587814301252365,-0.04344756156206131,0.020406663417816162,0.043269287794828415,0.04565715789794922,-0.07456950843334198,0.03894457593560219,0.03974293917417526,0.058624014258384705,0.04648742079734802,0.02159239910542965]]}' + headers: + content-length: + - '8050' + content-type: + - application/json; charset=utf-8 + content-version: + - '2023-11-22' + date: + - Thu, 08 Feb 2024 14:04:41 GMT + keep-alive: + - timeout=72 + transfer-encoding: + - chunked + vary: + - accept-encoding + status: + code: 200 + message: OK +version: 1 diff --git a/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml new file mode 100644 index 00000000..1019613a --- /dev/null +++ b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml @@ -0,0 +1,111 @@ +interactions: +- request: + body: '' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + method: GET + uri: https://api.com/v2/text/embeddings/limits?version=2023-11-22 + response: + body: + string: '{"result":{"concurrency":{"limit":20,"remaining":20}}}' + headers: + cache-control: + - private + content-length: + - '54' + content-type: + - application/json; charset=utf-8 + content-version: + - '2023-11-22' + date: + - Thu, 08 Feb 2024 13:48:59 GMT + keep-alive: + - timeout=72 + set-cookie: + - 2eef5f4c257f6bca76e8da5586743beb=10e3bbcca8eae261e1eab9f14bcfb604; path=/; + HttpOnly; Secure; SameSite=None + vary: + - accept-encoding + status: + code: 200 + message: OK +- request: + body: '{"input": "First document", "model_id": "sentence-transformers/all-minilm-l6-v2", + "parameters": {"truncate_input_tokens": true}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '128' + content-type: + - application/json + method: POST + uri: https://api.com/v2/text/embeddings?version=2023-11-22 + response: + body: + string: '{"results":[[-0.07611478865146637,0.053286295384168625,0.011630389839410782,0.02703363634645939,-0.020964650437235832,-0.0017028871225193143,-0.02678748220205307,0.04056990146636963,0.04763270914554596,0.014788346365094185,0.04339343681931496,0.07593807578086853,0.0012843846343457699,0.00855201669037342,-0.0745568573474884,0.006555262487381697,-0.018848007544875145,-0.0644168108701706,0.03469758480787277,0.06712162494659424,0.0019739633426070213,0.060395244508981705,-0.005633489694446325,-0.027578119188547134,0.017701366916298866,0.028293320909142494,-0.11111179739236832,0.0075835236348211765,0.04227476194500923,-0.11056772619485855,0.0488184429705143,-0.05738593265414238,0.09648681432008743,0.044576071202754974,0.05660592019557953,0.009718634188175201,0.05797331780195236,-0.05440652742981911,0.036171287298202515,0.06543885916471481,-0.0022649781312793493,-0.08455892652273178,-0.029309898614883423,-0.015962744131684303,0.01313595287501812,-0.0004417791496962309,0.010334083810448647,0.0992552787065506,0.049126140773296356,0.02620592713356018,-0.02891329675912857,-0.03497561067342758,-0.02703670598566532,0.018123308196663857,0.057025808840990067,0.04752988740801811,-0.015956465154886246,0.06666743755340576,-0.056488726288080215,-0.06855082511901855,-0.05432615429162979,0.007647170685231686,-0.07357412576675415,0.03254223242402077,0.12377501279115677,0.028229698538780212,-0.004680067300796509,-0.00489498907700181,-0.04303337633609772,0.04150404408574104,-0.016368089243769646,0.0326528362929821,0.01849137246608734,0.030499456450343132,-0.015698838979005814,-0.07027243822813034,0.006100979633629322,0.0350557342171669,0.04501684010028839,-0.09169086068868637,0.004178141709417105,-0.022076696157455444,0.030321242287755013,0.024455692619085312,-0.11486124247312546,-0.0025273847859352827,-0.019262436777353287,0.022405091673135757,-0.005074415821582079,-0.015164084732532501,0.007719150744378567,-0.04771441966295242,0.05346492677927017,0.044002655893564224,-0.07894007861614227,0.045628711581230164,-0.00509940879419446,-0.01611689105629921,0.10104403644800186,0.15719443559646606,0.0031960131600499153,0.023737283423542976,0.1408112645149231,0.06969963759183884,0.006006232928484678,-0.08215378224849701,-0.051280129700899124,-0.04340577498078346,-0.026042301207780838,-0.04198248311877251,-0.008222031407058239,-0.04777222126722336,-0.10308240354061127,0.013990739360451698,-0.04264850914478302,0.002755384426563978,0.031768303364515305,0.05728062614798546,0.02951638586819172,0.0031079030595719814,0.024997133761644363,0.010455358773469925,-0.02089112438261509,-0.030192524194717407,-0.1047578677535057,-0.0630701556801796,0.05935980752110481,-3.689310325890858e-33,0.018732953816652298,-0.00007181623368524015,-0.07653958350419998,0.08046584576368332,-0.03137052059173584,0.016886163502931595,0.01065531000494957,0.026087334379553795,-0.13024289906024933,0.021207323297858238,0.010911231860518456,0.010168812237679958,0.0050546638667583466,0.006856250576674938,-0.08671053498983383,0.034387242048978806,-0.0887618139386177,0.03489241376519203,-0.047422315925359726,-0.05843717232346535,-0.05840500816702843,-0.010714150965213776,0.05513739958405495,-0.004156034905463457,0.01576279290020466,0.033243175595998764,0.011517849750816822,-0.08593405783176422,-0.03567929565906525,-0.0134825948625803,-0.057146165519952774,-0.01318218745291233,-0.026386702433228493,-0.0748041421175003,-0.023853925988078117,0.10301246494054794,0.04259973391890526,-0.044059813022613525,-0.014225777238607407,-0.03309769555926323,-0.028882132843136787,0.05972263216972351,0.054447341710329056,-0.03648264333605766,-0.025904260575771332,0.032167814671993256,-0.05100731551647186,0.07830699533224106,0.12723541259765625,0.018466690555214882,-0.0486471951007843,0.017355622723698616,-0.045035943388938904,0.03343017399311066,-0.0060922387056052685,-0.03044017218053341,-0.0040046460926532745,-0.008548876270651817,-0.008502165786921978,-0.030725140124559402,0.018580880016088486,0.06358906626701355,-0.024671146646142006,0.06750013679265976,-0.013552273623645306,-0.048398926854133606,-0.031233591958880424,0.018136467784643173,0.09101729094982147,-0.08552554994821548,-0.10574617236852646,-0.03554975241422653,0.023665929213166237,-0.09693295508623123,0.0900241956114769,0.0746527761220932,0.07275242358446121,-0.056399762630462646,-0.060314252972602844,-0.04025064781308174,-0.09734656661748886,-0.0005814050091430545,-0.013118426315486431,0.0647367462515831,-0.029480375349521637,0.09067367017269135,0.08380653709173203,-0.0018367903539910913,-0.006479343399405479,-0.007722273003309965,0.00932267401367426,0.03696272149682045,0.05342373996973038,-0.018204154446721077,0.08472762256860733,2.457529632131353e-33,0.006457834504544735,-0.055136796087026596,0.006394235882908106,0.04257359355688095,0.040424179285764694,0.030993953347206116,0.04744979366660118,0.02580159716308117,-0.010715586133301258,0.05614149942994118,0.07997943460941315,-0.012969913892447948,0.04406759887933731,-0.007733637932687998,0.01743471808731556,-0.024840814992785454,0.05109093710780144,-0.032808005809783936,0.03324175998568535,0.051604095846414566,-0.01836460828781128,-0.04881879314780235,-0.03895711898803711,0.06022154167294502,0.018858902156352997,0.01531311683356762,0.07880528271198273,-0.057667337357997894,-0.07734047621488571,-0.029612315818667412,-0.017924057319760323,-0.007808707654476166,-0.005763333290815353,0.07769045233726501,0.013317340053617954,-0.004131403751671314,-0.006829418241977692,-0.006058764178305864,0.020915012806653976,0.08015037328004837,0.06981609016656876,0.022621115669608116,-0.011095408350229263,-0.0010161078535020351,-0.047971002757549286,-0.06444189697504044,0.024511374533176422,0.02654142677783966,0.0009377644746564329,-0.042746491730213165,-0.03214338794350624,0.04735948517918587,0.022207461297512054,-0.14833605289459229,-0.012862633913755417,0.001661280868574977,-0.01834721677005291,-0.03249534219503403,0.0025780177675187588,0.07029899209737778,0.03629522770643234,0.08275070786476135,-0.06574779003858566,0.047293007373809814,-0.042443957179784775,-0.09515932947397232,-0.07330439239740372,0.05922340974211693,-0.0368594266474247,0.01581326127052307,-0.041240546852350235,0.02242102473974228,0.007569017820060253,-0.05807928368449211,0.05175217613577843,-0.017240695655345917,0.033588897436857224,-0.03882821276783943,-0.0032414624001830816,-0.07458201050758362,0.060188066214323044,0.05588485673069954,-0.07107950747013092,0.007625015452504158,0.09588012099266052,-0.0013106901897117496,0.05206756293773651,-0.11739741265773773,-0.032384924590587616,0.028821000829339027,-0.10565381497144699,0.040773868560791016,0.05618344619870186,0.06782420724630356,-0.08374620229005814,-1.2194121268294111e-8,-0.10979562252759933,0.014475689269602299,0.006779859308153391,0.018453504890203476,0.018357109278440475,0.04331545531749725,0.025635426864027977,0.023104725405573845,-0.08535239100456238,-0.10341934859752655,0.01501353457570076,0.03674392029643059,0.04435298219323158,-0.05127701163291931,-0.011615892872214317,-0.01459368597716093,0.0279467161744833,-0.10243868082761765,-0.05912458151578903,0.004412763752043247,0.022360317409038544,0.013892627321183681,0.0198416318744421,-0.0965014323592186,0.028432726860046387,0.003722630674019456,0.04970787465572357,0.019583890214562416,0.03422922641038895,-0.0013914067531004548,-0.027568506076931953,0.10341252386569977,-0.01579933986067772,0.024184221401810646,0.05524905025959015,0.03001238778233528,0.0234824325889349,0.06176608055830002,-0.03289225324988365,-0.06597874313592911,0.021079285070300102,0.026373285800218582,-0.04654461145401001,-0.01519238855689764,-0.032833341509103775,0.003992781043052673,0.0065460167825222015,-0.05169658362865448,0.021103447303175926,-0.10728193819522858,-0.02706659585237503,-0.06279673427343369,0.1034388542175293,0.056007225066423416,-0.02288142964243889,0.03683370351791382,0.06275887787342072,0.013564606197178364,-0.05565091222524643,0.03672057390213013,0.03527916222810745,0.066694475710392,0.09273931384086609,0.027666479349136353]]}' + headers: + content-length: + - '8058' + content-type: + - application/json; charset=utf-8 + content-version: + - '2023-11-22' + date: + - Thu, 08 Feb 2024 13:48:59 GMT + keep-alive: + - timeout=72 + transfer-encoding: + - chunked + vary: + - accept-encoding + status: + code: 200 + message: OK +- request: + body: '{"input": "Second document", "model_id": "sentence-transformers/all-minilm-l6-v2", + "parameters": {"truncate_input_tokens": true}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '129' + content-type: + - application/json + method: POST + uri: https://api.com/v2/text/embeddings?version=2023-11-22 + response: + body: + string: '{"results":[[-0.08104048669338226,0.0376598946750164,0.0020215765107423067,0.026549046859145164,-0.007370871026068926,0.011257929727435112,-0.02854747697710991,0.0631430447101593,0.03664477914571762,-0.03365236148238182,0.06271136552095413,0.06864234060049057,0.02026943676173687,-0.030009351670742035,-0.03508774936199188,0.0023406269028782845,-0.08050798624753952,-0.016443617641925812,0.0018381527625024319,0.062218863517045975,0.03620733320713043,0.009204098023474216,-0.015127652324736118,-0.025030456483364105,0.010739271529018879,0.06100801005959511,-0.11221961677074432,0.05921312794089317,0.0026483291294425726,-0.10150498896837234,0.05173927918076515,0.023861411958932877,-0.033372968435287476,-0.001777892466634512,0.021250296384096146,-0.058816056698560715,0.018938081339001656,-0.00041302599129267037,0.043738409876823425,-0.004523207433521748,-0.011076416820287704,-0.06403172016143799,-0.05241154134273529,-0.04321694001555443,-0.017592426389455795,0.011236284859478474,-0.04660659283399582,0.07122266292572021,0.044390350580215454,0.0497487410902977,-0.07200715690851212,-0.04020734876394272,-0.06696461886167526,0.029645787551999092,0.0457756444811821,0.11751767247915268,-0.012403221800923347,0.056834425777196884,-0.04663803055882454,-0.06756290793418884,0.0081444401293993,0.05049721896648407,-0.10587693005800247,0.0636945515871048,0.1167907789349556,-0.012800364755094051,-0.016207052394747734,-0.004366802517324686,-0.11997950077056885,0.026186034083366394,0.014950335957109928,0.03286281228065491,0.07058639079332352,-0.03315839171409607,0.06088574603199959,-0.10461709648370743,-0.02542501501739025,0.025162100791931152,0.014831612817943096,-0.15255455672740936,0.02012590691447258,0.010713051073253155,0.014526871033012867,-0.04263008385896683,-0.08039771765470505,-0.0731976181268692,-0.00012689009599853307,0.0034983744844794273,-0.022338520735502243,-0.006795673631131649,0.0667412057518959,-0.055777598172426224,0.11562377214431763,0.034403957426548004,-0.06897035986185074,0.0018935615662485361,-0.002328364411368966,0.044492337852716446,0.11536921560764313,0.14253538846969604,-0.017649555578827858,0.00923872459679842,0.08490171283483505,-0.00015613385767210275,0.04251273721456528,-0.09667838364839554,-0.0794786885380745,-0.06583905220031738,0.037441436201334,-0.061086781322956085,0.02558034099638462,0.014064162038266659,-0.09667671471834183,0.026026040315628052,0.01610356941819191,0.06567886471748352,0.08934935927391052,0.05728030577301979,0.03825286403298378,-0.004467225167900324,-0.03369850665330887,-0.011704717762768269,0.015539701096713543,-0.060387320816516876,-0.06724131852388382,-0.10934903472661972,0.07420971244573593,-3.517199218489259e-33,-0.0010445192456245422,-0.04707157239317894,-0.09643448889255524,0.061116773635149,0.013092653825879097,0.048495952039957047,0.012917639687657356,0.05141948536038399,-0.11057691276073456,-0.02722175046801567,-0.013994114473462105,0.014674837701022625,0.007228940725326538,0.0286633912473917,-0.06814133375883102,0.002892223419621587,-0.013423651456832886,0.09643645584583282,0.0409725159406662,-0.012952705845236778,0.009580191224813461,0.06898463517427444,0.0559326708316803,0.013785253278911114,0.05769336596131325,0.07820483297109604,0.04037198796868324,-0.06474727392196655,-0.045225292444229126,-0.002468392252922058,0.031135447323322296,-0.003809040179476142,0.016344018280506134,-0.030952706933021545,0.0029499472584575415,0.07108024507761002,0.020287014544010162,-0.027239948511123657,0.01470908336341381,-0.03253556415438652,-0.02227545529603958,0.023334389552474022,0.04677804186940193,-0.02296368218958378,0.01581859588623047,-0.012426641769707203,-0.009094944223761559,0.030619749799370766,0.12037337571382523,-0.01623496040701866,-0.020219366997480392,0.026702173054218292,-0.056418877094984055,0.020353808999061584,0.03395252302289009,-0.009550661779940128,-0.0622478723526001,0.022805457934737206,0.062185581773519516,0.001912244246341288,0.0438990518450737,0.08762244880199432,-0.025203194469213486,0.04708396643400192,-0.0020530272740870714,-0.003188277827575803,-0.08097474277019501,0.0046669184230268,0.09208934754133224,-0.03296930342912674,-0.10864203423261642,-0.07346982508897781,0.06713204085826874,-0.0824589729309082,0.08118608593940735,-0.015304031781852245,0.03169460967183113,-0.05894801393151283,-0.036768555641174316,-0.052348583936691284,-0.13439710438251495,-0.019004182890057564,-0.029325569048523903,0.0389697328209877,-0.03179283067584038,0.07625920325517654,0.029881611466407776,-0.0287876445800066,-0.03186427801847458,0.051656320691108704,-0.020476073026657104,0.06108597293496132,-0.03309909999370575,-0.03133446350693703,0.14522379636764526,1.9997601731684913e-33,-0.02454514242708683,-0.02168901264667511,-0.04479790851473808,0.04633335769176483,0.033506955951452255,0.039710696786642075,0.04203144833445549,-0.025988228619098663,-0.059948988258838654,0.04111616685986519,0.04170918092131615,-0.010260396637022495,0.08616381138563156,-0.01634076237678528,-0.023933496326208115,-0.03121459297835827,0.053647492080926895,0.009220186620950699,0.013340904377400875,0.0437324121594429,-0.02444729581475258,-0.03904367983341217,0.030613556504249573,0.01505532581359148,0.05224118381738663,-0.002882804023101926,0.011661234311759472,-0.0403885692358017,-0.05561037361621857,-0.03922025114297867,-0.017746806144714355,-0.024655308574438095,-0.04797155410051346,0.03661811724305153,0.046329937875270844,-0.03573417663574219,0.0007827881490811706,-0.032068975269794464,-0.04048440232872963,0.0529610775411129,0.06700104475021362,0.05207323282957077,0.025980258360505104,0.003316115355119109,-0.05455101653933525,0.021097073331475258,0.015427522361278534,0.012054149992763996,0.0036217537708580494,-0.047605108469724655,-0.0024289479479193687,-0.023437868803739548,-0.016562774777412415,-0.09189718216657639,-0.06719052046537399,0.009031710214912891,0.015131138265132904,-0.018840065225958824,0.03931032866239548,0.04334070906043053,0.09708531945943832,0.04747961834073067,-0.06743241101503372,0.03180370107293129,0.017487889155745506,-0.08593153208494186,-0.07724672555923462,0.02401409111917019,-0.03559126332402229,0.07891940325498581,-0.11396396905183792,0.000023013402824290097,0.03012612834572792,-0.05196437984704971,0.05744696781039238,0.00019338278798386455,0.02741345390677452,-0.05575890466570854,-0.008595319464802742,-0.051599353551864624,0.06556955724954605,0.03357896953821182,-0.02616681344807148,0.022066596895456314,-0.0033872537314891815,-0.000342347368132323,0.020659705623984337,-0.03631681948900223,-0.03139133006334305,-0.01155750174075365,-0.06169193983078003,0.08895165473222733,-0.013894675299525261,0.012268207967281342,-0.03348323702812195,-1.2230937151969101e-8,-0.06460992246866226,-0.017255505546927452,0.036914072930812836,-0.036867473274469376,0.005709514953196049,-0.023781590163707733,-0.008763302117586136,0.0318242646753788,-0.06689297407865524,-0.06356614083051682,0.012779262848198414,0.01958351768553257,0.06348159909248352,-0.03766090050339699,-0.034059420228004456,-0.04260662943124771,0.07547925412654877,-0.10275367647409439,-0.04524024948477745,-0.03890029340982437,0.026446901261806488,-0.009833211079239845,0.03849169984459877,-0.009132797829806805,-0.034811895340681076,0.05913154408335686,0.04117123410105705,-0.015952495858073235,-0.013833853416144848,-0.051513154059648514,0.017751075327396393,0.06703679263591766,-0.03192468360066414,0.02006058767437935,0.0793311819434166,-0.004208602476865053,0.045045576989650726,0.07222826033830643,-0.04168285056948662,0.0011332163121551275,0.04194065183401108,-0.024034954607486725,-0.053235992789268494,0.05583583563566208,0.08766256272792816,-0.014262460172176361,0.0411042682826519,-0.10182081907987595,0.029361538589000702,-0.04424484819173813,0.01873917318880558,-0.0940093845129013,0.1568623185157776,0.050587814301252365,-0.04344756156206131,0.020406663417816162,0.043269287794828415,0.04565715789794922,-0.07456950843334198,0.03894457593560219,0.03974293917417526,0.058624014258384705,0.04648742079734802,0.02159239910542965]]}' + headers: + content-length: + - '8050' + content-type: + - application/json; charset=utf-8 + content-version: + - '2023-11-22' + date: + - Thu, 08 Feb 2024 13:48:59 GMT + keep-alive: + - timeout=72 + transfer-encoding: + - chunked + vary: + - accept-encoding + status: + code: 200 + message: OK +version: 1 diff --git a/tests/integration/extensions/test_llama_index_embeddings.py b/tests/integration/extensions/test_llama_index_embeddings.py new file mode 100644 index 00000000..7820e182 --- /dev/null +++ b/tests/integration/extensions/test_llama_index_embeddings.py @@ -0,0 +1,39 @@ +import pytest + +from genai.extensions.llama_index import LlamaIndexEmbeddingsInterface +from genai.schema import TextEmbeddingCreateEndpoint, TextEmbeddingParameters + + +@pytest.mark.integration +class TestLangChainEmbeddings: + def setup_method(self): + self.model_id = "sentence-transformers/all-minilm-l6-v2" + + @pytest.fixture + def parameters(self): + return TextEmbeddingParameters(truncate_input_tokens=True) + + @pytest.fixture + def documents(self) -> list[str]: + return ["First document", "Second document"] + + @pytest.mark.vcr + def test_embedding(self, client, parameters, documents, get_vcr_responses_of): + model = LlamaIndexEmbeddingsInterface(model_id=self.model_id, parameters=parameters, client=client) + + results = model.get_text_embedding_batch(documents) + responses = get_vcr_responses_of(TextEmbeddingCreateEndpoint) + + for response, result in zip(responses, results): + assert result == response["results"][0] + + @pytest.mark.asyncio + @pytest.mark.vcr + async def test_async_embedding(self, client, parameters, documents, get_vcr_responses_of): + model = LlamaIndexEmbeddingsInterface(model_id=self.model_id, parameters=parameters, client=client) + + results = await model.aget_text_embedding_batch(documents) + responses = get_vcr_responses_of(TextEmbeddingCreateEndpoint) + + for response, result in zip(responses, results): + assert result == response["results"][0] From 746b94ec575b4c3276752a08c393fa520b2e8a63 Mon Sep 17 00:00:00 2001 From: David Kristek Date: Thu, 8 Feb 2024 17:17:13 +0100 Subject: [PATCH 2/2] fix: correct naming --- examples/extensions/llama_index/llama_index_embedding.py | 6 +++--- src/genai/extensions/llama_index/__init__.py | 4 ++-- src/genai/extensions/llama_index/embeddings.py | 2 +- ...=> TestLlamaIndexEmbeddings.test_async_embedding.yaml} | 6 +++--- ....yaml => TestLlamaIndexEmbeddings.test_embedding.yaml} | 6 +++--- .../integration/extensions/test_llama_index_embeddings.py | 8 ++++---- 6 files changed, 16 insertions(+), 16 deletions(-) rename tests/integration/extensions/cassettes/test_llama_index_embeddings/{TestLangChainEmbeddings.test_async_embedding.yaml => TestLlamaIndexEmbeddings.test_async_embedding.yaml} (99%) rename tests/integration/extensions/cassettes/test_llama_index_embeddings/{TestLangChainEmbeddings.test_embedding.yaml => TestLlamaIndexEmbeddings.test_embedding.yaml} (99%) diff --git a/examples/extensions/llama_index/llama_index_embedding.py b/examples/extensions/llama_index/llama_index_embedding.py index d33082eb..64f5c6ed 100644 --- a/examples/extensions/llama_index/llama_index_embedding.py +++ b/examples/extensions/llama_index/llama_index_embedding.py @@ -1,10 +1,10 @@ """ -Llamaindex Embeddings +LlamaIndex Embeddings """ from dotenv import load_dotenv from genai import Client, Credentials -from genai.extensions.llama_index import LlamaIndexEmbeddingsInterface +from genai.extensions.llama_index import IBMGenAILlamaIndexEmbedding from genai.schema import TextEmbeddingParameters # make sure you have a .env file under genai root with @@ -21,7 +21,7 @@ def heading(text: str) -> str: print(heading("LlamaIndex Embeddings")) client = Client(credentials=Credentials.from_env()) -embeddings = LlamaIndexEmbeddingsInterface( +embeddings = IBMGenAILlamaIndexEmbedding( client=client, model_id="sentence-transformers/all-minilm-l6-v2", parameters=TextEmbeddingParameters(truncate_input_tokens=True), diff --git a/src/genai/extensions/llama_index/__init__.py b/src/genai/extensions/llama_index/__init__.py index 80a68d68..19cc0bca 100644 --- a/src/genai/extensions/llama_index/__init__.py +++ b/src/genai/extensions/llama_index/__init__.py @@ -1,6 +1,6 @@ """Extension for LLamaIndex library""" -from genai.extensions.llama_index.embeddings import LlamaIndexEmbeddingsInterface +from genai.extensions.llama_index.embeddings import IBMGenAILlamaIndexEmbedding from genai.extensions.llama_index.llm import IBMGenAILlamaIndex -__all__ = ["IBMGenAILlamaIndex", "LlamaIndexEmbeddingsInterface"] +__all__ = ["IBMGenAILlamaIndex", "IBMGenAILlamaIndexEmbedding"] diff --git a/src/genai/extensions/llama_index/embeddings.py b/src/genai/extensions/llama_index/embeddings.py index 2e5bdd6c..89829d2f 100644 --- a/src/genai/extensions/llama_index/embeddings.py +++ b/src/genai/extensions/llama_index/embeddings.py @@ -12,7 +12,7 @@ raise ImportError("Could not import llamaindex: Please install ibm-generative-ai[llama-index] extension.") # noqa: B904 -class LlamaIndexEmbeddingsInterface(BaseEmbedding): +class IBMGenAILlamaIndexEmbedding(BaseEmbedding): client: Client model_id: str parameters: Optional[ModelLike[TextEmbeddingParameters]] = None diff --git a/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLlamaIndexEmbeddings.test_async_embedding.yaml similarity index 99% rename from tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml rename to tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLlamaIndexEmbeddings.test_async_embedding.yaml index 2b189c1b..55056e72 100644 --- a/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_async_embedding.yaml +++ b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLlamaIndexEmbeddings.test_async_embedding.yaml @@ -23,7 +23,7 @@ interactions: content-version: - '2023-11-22' date: - - Thu, 08 Feb 2024 14:04:40 GMT + - Thu, 08 Feb 2024 16:15:16 GMT keep-alive: - timeout=72 set-cookie: @@ -61,7 +61,7 @@ interactions: content-version: - '2023-11-22' date: - - Thu, 08 Feb 2024 14:04:41 GMT + - Thu, 08 Feb 2024 16:15:16 GMT keep-alive: - timeout=72 transfer-encoding: @@ -98,7 +98,7 @@ interactions: content-version: - '2023-11-22' date: - - Thu, 08 Feb 2024 14:04:41 GMT + - Thu, 08 Feb 2024 16:15:16 GMT keep-alive: - timeout=72 transfer-encoding: diff --git a/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLlamaIndexEmbeddings.test_embedding.yaml similarity index 99% rename from tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml rename to tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLlamaIndexEmbeddings.test_embedding.yaml index 1019613a..de218c00 100644 --- a/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLangChainEmbeddings.test_embedding.yaml +++ b/tests/integration/extensions/cassettes/test_llama_index_embeddings/TestLlamaIndexEmbeddings.test_embedding.yaml @@ -23,7 +23,7 @@ interactions: content-version: - '2023-11-22' date: - - Thu, 08 Feb 2024 13:48:59 GMT + - Thu, 08 Feb 2024 16:15:15 GMT keep-alive: - timeout=72 set-cookie: @@ -61,7 +61,7 @@ interactions: content-version: - '2023-11-22' date: - - Thu, 08 Feb 2024 13:48:59 GMT + - Thu, 08 Feb 2024 16:15:15 GMT keep-alive: - timeout=72 transfer-encoding: @@ -98,7 +98,7 @@ interactions: content-version: - '2023-11-22' date: - - Thu, 08 Feb 2024 13:48:59 GMT + - Thu, 08 Feb 2024 16:15:15 GMT keep-alive: - timeout=72 transfer-encoding: diff --git a/tests/integration/extensions/test_llama_index_embeddings.py b/tests/integration/extensions/test_llama_index_embeddings.py index 7820e182..e87ffcbc 100644 --- a/tests/integration/extensions/test_llama_index_embeddings.py +++ b/tests/integration/extensions/test_llama_index_embeddings.py @@ -1,11 +1,11 @@ import pytest -from genai.extensions.llama_index import LlamaIndexEmbeddingsInterface +from genai.extensions.llama_index import IBMGenAILlamaIndexEmbedding from genai.schema import TextEmbeddingCreateEndpoint, TextEmbeddingParameters @pytest.mark.integration -class TestLangChainEmbeddings: +class TestLlamaIndexEmbeddings: def setup_method(self): self.model_id = "sentence-transformers/all-minilm-l6-v2" @@ -19,7 +19,7 @@ def documents(self) -> list[str]: @pytest.mark.vcr def test_embedding(self, client, parameters, documents, get_vcr_responses_of): - model = LlamaIndexEmbeddingsInterface(model_id=self.model_id, parameters=parameters, client=client) + model = IBMGenAILlamaIndexEmbedding(model_id=self.model_id, parameters=parameters, client=client) results = model.get_text_embedding_batch(documents) responses = get_vcr_responses_of(TextEmbeddingCreateEndpoint) @@ -30,7 +30,7 @@ def test_embedding(self, client, parameters, documents, get_vcr_responses_of): @pytest.mark.asyncio @pytest.mark.vcr async def test_async_embedding(self, client, parameters, documents, get_vcr_responses_of): - model = LlamaIndexEmbeddingsInterface(model_id=self.model_id, parameters=parameters, client=client) + model = IBMGenAILlamaIndexEmbedding(model_id=self.model_id, parameters=parameters, client=client) results = await model.aget_text_embedding_batch(documents) responses = get_vcr_responses_of(TextEmbeddingCreateEndpoint)