Skip to content

Commit 37d1da6

Browse files
Unprototyping Cognition (#177)
* creates alembic migration for strategy step refactoring * updates submodule
1 parent 7358312 commit 37d1da6

File tree

2 files changed

+347
-1
lines changed

2 files changed

+347
-1
lines changed
Lines changed: 346 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,346 @@
1+
"""refactor cognition strategy step
2+
3+
Revision ID: 3d0e01981f06
4+
Revises: 491ea68a7baf
5+
Create Date: 2023-12-05 10:38:21.403038
6+
7+
"""
8+
from alembic import op
9+
import sqlalchemy as sa
10+
from sqlalchemy.dialects import postgresql
11+
12+
# revision identifiers, used by Alembic.
13+
revision = "3d0e01981f06"
14+
down_revision = "491ea68a7baf"
15+
branch_labels = None
16+
depends_on = None
17+
18+
19+
def upgrade():
20+
# ### commands auto generated by Alembic - please adjust! ###
21+
op.drop_index(
22+
"ix_cognition_python_step_created_by",
23+
table_name="python_step",
24+
schema="cognition",
25+
)
26+
op.drop_index(
27+
"ix_cognition_python_step_project_id",
28+
table_name="python_step",
29+
schema="cognition",
30+
)
31+
op.drop_index(
32+
"ix_cognition_python_step_strategy_step_id",
33+
table_name="python_step",
34+
schema="cognition",
35+
)
36+
op.drop_table("python_step", schema="cognition")
37+
op.drop_index(
38+
"ix_cognition_retriever_part_created_by",
39+
table_name="retriever_part",
40+
schema="cognition",
41+
)
42+
op.drop_index(
43+
"ix_cognition_retriever_part_project_id",
44+
table_name="retriever_part",
45+
schema="cognition",
46+
)
47+
op.drop_index(
48+
"ix_cognition_retriever_part_retriever_id",
49+
table_name="retriever_part",
50+
schema="cognition",
51+
)
52+
op.drop_table("retriever_part", schema="cognition")
53+
op.drop_index(
54+
"ix_cognition_retriever_created_by", table_name="retriever", schema="cognition"
55+
)
56+
op.drop_index(
57+
"ix_cognition_retriever_project_id", table_name="retriever", schema="cognition"
58+
)
59+
op.drop_index(
60+
"ix_cognition_retriever_strategy_step_id",
61+
table_name="retriever",
62+
schema="cognition",
63+
)
64+
op.drop_table("retriever", schema="cognition")
65+
op.drop_index(
66+
"ix_cognition_llm_step_created_by", table_name="llm_step", schema="cognition"
67+
)
68+
op.drop_index(
69+
"ix_cognition_llm_step_project_id", table_name="llm_step", schema="cognition"
70+
)
71+
op.drop_index(
72+
"ix_cognition_llm_step_strategy_step_id",
73+
table_name="llm_step",
74+
schema="cognition",
75+
)
76+
op.drop_table("llm_step", schema="cognition")
77+
op.add_column(
78+
"strategy_step",
79+
sa.Column("step_type", sa.String(), nullable=True),
80+
schema="cognition",
81+
)
82+
op.add_column(
83+
"strategy_step",
84+
sa.Column("position", sa.Integer(), nullable=True),
85+
schema="cognition",
86+
)
87+
op.add_column(
88+
"strategy_step",
89+
sa.Column("config", sa.JSON(), nullable=True),
90+
schema="cognition",
91+
)
92+
op.drop_column("strategy_step", "strategy_step_position", schema="cognition")
93+
op.drop_column("strategy_step", "strategy_step_type", schema="cognition")
94+
# ### end Alembic commands ###
95+
96+
97+
def downgrade():
98+
# ### commands auto generated by Alembic - please adjust! ###
99+
op.add_column(
100+
"strategy_step",
101+
sa.Column(
102+
"strategy_step_type", sa.VARCHAR(), autoincrement=False, nullable=True
103+
),
104+
schema="cognition",
105+
)
106+
op.add_column(
107+
"strategy_step",
108+
sa.Column(
109+
"strategy_step_position", sa.INTEGER(), autoincrement=False, nullable=True
110+
),
111+
schema="cognition",
112+
)
113+
op.drop_column("strategy_step", "config", schema="cognition")
114+
op.drop_column("strategy_step", "position", schema="cognition")
115+
op.drop_column("strategy_step", "step_type", schema="cognition")
116+
op.create_table(
117+
"llm_step",
118+
sa.Column("id", postgresql.UUID(), autoincrement=False, nullable=False),
119+
sa.Column("project_id", postgresql.UUID(), autoincrement=False, nullable=True),
120+
sa.Column(
121+
"strategy_step_id", postgresql.UUID(), autoincrement=False, nullable=True
122+
),
123+
sa.Column("created_by", postgresql.UUID(), autoincrement=False, nullable=True),
124+
sa.Column(
125+
"created_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=True
126+
),
127+
sa.Column("llm_identifier", sa.VARCHAR(), autoincrement=False, nullable=True),
128+
sa.Column(
129+
"llm_config",
130+
postgresql.JSON(astext_type=sa.Text()),
131+
autoincrement=False,
132+
nullable=True,
133+
),
134+
sa.Column("template_prompt", sa.VARCHAR(), autoincrement=False, nullable=True),
135+
sa.Column("question_prompt", sa.VARCHAR(), autoincrement=False, nullable=True),
136+
sa.ForeignKeyConstraint(
137+
["created_by"],
138+
["user.id"],
139+
name="llm_step_created_by_fkey",
140+
ondelete="SET NULL",
141+
),
142+
sa.ForeignKeyConstraint(
143+
["project_id"],
144+
["cognition.project.id"],
145+
name="llm_step_project_id_fkey",
146+
ondelete="CASCADE",
147+
),
148+
sa.ForeignKeyConstraint(
149+
["strategy_step_id"],
150+
["cognition.strategy_step.id"],
151+
name="llm_step_strategy_step_id_fkey",
152+
ondelete="CASCADE",
153+
),
154+
sa.PrimaryKeyConstraint("id", name="llm_step_pkey"),
155+
schema="cognition",
156+
)
157+
op.create_index(
158+
"ix_cognition_llm_step_strategy_step_id",
159+
"llm_step",
160+
["strategy_step_id"],
161+
unique=False,
162+
schema="cognition",
163+
)
164+
op.create_index(
165+
"ix_cognition_llm_step_project_id",
166+
"llm_step",
167+
["project_id"],
168+
unique=False,
169+
schema="cognition",
170+
)
171+
op.create_index(
172+
"ix_cognition_llm_step_created_by",
173+
"llm_step",
174+
["created_by"],
175+
unique=False,
176+
schema="cognition",
177+
)
178+
op.create_table(
179+
"retriever",
180+
sa.Column("id", postgresql.UUID(), autoincrement=False, nullable=False),
181+
sa.Column("project_id", postgresql.UUID(), autoincrement=False, nullable=True),
182+
sa.Column(
183+
"strategy_step_id", postgresql.UUID(), autoincrement=False, nullable=True
184+
),
185+
sa.Column("created_by", postgresql.UUID(), autoincrement=False, nullable=True),
186+
sa.Column(
187+
"created_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=True
188+
),
189+
sa.Column(
190+
"search_input_field", sa.VARCHAR(), autoincrement=False, nullable=True
191+
),
192+
sa.Column("meta_data", sa.VARCHAR(), autoincrement=False, nullable=True),
193+
sa.ForeignKeyConstraint(
194+
["created_by"],
195+
["user.id"],
196+
name="retriever_created_by_fkey",
197+
ondelete="SET NULL",
198+
),
199+
sa.ForeignKeyConstraint(
200+
["project_id"],
201+
["cognition.project.id"],
202+
name="retriever_project_id_fkey",
203+
ondelete="CASCADE",
204+
),
205+
sa.ForeignKeyConstraint(
206+
["strategy_step_id"],
207+
["cognition.strategy_step.id"],
208+
name="retriever_strategy_step_id_fkey",
209+
ondelete="CASCADE",
210+
),
211+
sa.PrimaryKeyConstraint("id", name="retriever_pkey"),
212+
schema="cognition",
213+
postgresql_ignore_search_path=False,
214+
)
215+
op.create_index(
216+
"ix_cognition_retriever_strategy_step_id",
217+
"retriever",
218+
["strategy_step_id"],
219+
unique=False,
220+
schema="cognition",
221+
)
222+
op.create_index(
223+
"ix_cognition_retriever_project_id",
224+
"retriever",
225+
["project_id"],
226+
unique=False,
227+
schema="cognition",
228+
)
229+
op.create_index(
230+
"ix_cognition_retriever_created_by",
231+
"retriever",
232+
["created_by"],
233+
unique=False,
234+
schema="cognition",
235+
)
236+
op.create_table(
237+
"retriever_part",
238+
sa.Column("id", postgresql.UUID(), autoincrement=False, nullable=False),
239+
sa.Column("project_id", postgresql.UUID(), autoincrement=False, nullable=True),
240+
sa.Column(
241+
"retriever_id", postgresql.UUID(), autoincrement=False, nullable=True
242+
),
243+
sa.Column("created_by", postgresql.UUID(), autoincrement=False, nullable=True),
244+
sa.Column(
245+
"created_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=True
246+
),
247+
sa.Column("embedding_name", sa.VARCHAR(), autoincrement=False, nullable=True),
248+
sa.Column("number_records", sa.INTEGER(), autoincrement=False, nullable=True),
249+
sa.Column("enabled", sa.BOOLEAN(), autoincrement=False, nullable=True),
250+
sa.ForeignKeyConstraint(
251+
["created_by"],
252+
["user.id"],
253+
name="retriever_part_created_by_fkey",
254+
ondelete="SET NULL",
255+
),
256+
sa.ForeignKeyConstraint(
257+
["project_id"],
258+
["cognition.project.id"],
259+
name="retriever_part_project_id_fkey",
260+
ondelete="CASCADE",
261+
),
262+
sa.ForeignKeyConstraint(
263+
["retriever_id"],
264+
["cognition.retriever.id"],
265+
name="retriever_part_retriever_id_fkey",
266+
ondelete="CASCADE",
267+
),
268+
sa.PrimaryKeyConstraint("id", name="retriever_part_pkey"),
269+
schema="cognition",
270+
)
271+
op.create_index(
272+
"ix_cognition_retriever_part_retriever_id",
273+
"retriever_part",
274+
["retriever_id"],
275+
unique=False,
276+
schema="cognition",
277+
)
278+
op.create_index(
279+
"ix_cognition_retriever_part_project_id",
280+
"retriever_part",
281+
["project_id"],
282+
unique=False,
283+
schema="cognition",
284+
)
285+
op.create_index(
286+
"ix_cognition_retriever_part_created_by",
287+
"retriever_part",
288+
["created_by"],
289+
unique=False,
290+
schema="cognition",
291+
)
292+
op.create_table(
293+
"python_step",
294+
sa.Column("id", postgresql.UUID(), autoincrement=False, nullable=False),
295+
sa.Column("project_id", postgresql.UUID(), autoincrement=False, nullable=True),
296+
sa.Column(
297+
"strategy_step_id", postgresql.UUID(), autoincrement=False, nullable=True
298+
),
299+
sa.Column("created_by", postgresql.UUID(), autoincrement=False, nullable=True),
300+
sa.Column(
301+
"created_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=True
302+
),
303+
sa.Column("source_code", sa.VARCHAR(), autoincrement=False, nullable=True),
304+
sa.ForeignKeyConstraint(
305+
["created_by"],
306+
["user.id"],
307+
name="python_step_created_by_fkey",
308+
ondelete="SET NULL",
309+
),
310+
sa.ForeignKeyConstraint(
311+
["project_id"],
312+
["cognition.project.id"],
313+
name="python_step_project_id_fkey",
314+
ondelete="CASCADE",
315+
),
316+
sa.ForeignKeyConstraint(
317+
["strategy_step_id"],
318+
["cognition.strategy_step.id"],
319+
name="python_step_strategy_step_id_fkey",
320+
ondelete="CASCADE",
321+
),
322+
sa.PrimaryKeyConstraint("id", name="python_step_pkey"),
323+
schema="cognition",
324+
)
325+
op.create_index(
326+
"ix_cognition_python_step_strategy_step_id",
327+
"python_step",
328+
["strategy_step_id"],
329+
unique=False,
330+
schema="cognition",
331+
)
332+
op.create_index(
333+
"ix_cognition_python_step_project_id",
334+
"python_step",
335+
["project_id"],
336+
unique=False,
337+
schema="cognition",
338+
)
339+
op.create_index(
340+
"ix_cognition_python_step_created_by",
341+
"python_step",
342+
["created_by"],
343+
unique=False,
344+
schema="cognition",
345+
)
346+
# ### end Alembic commands ###

0 commit comments

Comments
 (0)