1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
|
"""Example of using LLMs with PyOMOP
git clone https://github.com/dermatologist/pyomop.git@develop
cd pyomop
pip install -e .[llm]
"""
from pyomop import CdmEngineFactory, Cohort, metadata, CdmLLMQuery, CDMDatabase
import datetime
import asyncio
# Import any LLMs that llama_index supports and you have access to
# Require OpenAI API key to use OpenAI LLMs
from llama_index.llms import Vertex
async def main():
# Create a sqllite database by default
# You can also connect to an existing CDM database using the CdmEngineFactory
cdm = CdmEngineFactory()
# Postgres example below (db='mysql' also supported)
# cdm = CdmEngineFactory(db='pgsql', host='', port=5432,
# user='', pw='',
# name='', schema='cdm6')
engine = cdm.engine
# Create Tables if required
await cdm.init_models(metadata)
async with cdm.session() as session:
async with session.begin():
# Adding a cohort just for the example (not required if you have a OMOP CDM database)
session.add(Cohort(cohort_definition_id=2, subject_id=100,
cohort_end_date=datetime.datetime.now(),
cohort_start_date=datetime.datetime.now()))
await session.commit()
# Use any LLM that llama_index supports
llm = Vertex(
model="chat-bison",
)
# Include tables that you want to query
sql_database = CDMDatabase(engine, include_tables=[
"cohort",
])
query_engine = CdmLLMQuery(sql_database, llm=llm)
# Try any complex query.
response = query_engine.query("Show each in table cohort with a subject id of 100?")
print(response)
"""
| cohort_id | subject_id | cohort_name |
|---|---|---|
| 1 | 100 | Math |
"""
# Close session
await session.close()
await engine.dispose()
# Run the main function
asyncio.run(main())
|
