update python example

examples/python/client.py

@@ -3,10 +3,8 @@ import requests
 
 # NOTE: ollama must be running for this to work, start the ollama app or run `ollama serve`
 model = 'llama2' # TODO: update this for whatever model you wish to use
-context = [] # the context stores a conversation history, you can use this to make the model more context aware
 
-def generate(prompt):
-    global context
+def generate(prompt, context):
     r = requests.post('http://localhost:11434/api/generate',
                       json={
                           'model': model,
@@ -26,14 +24,14 @@ def generate(prompt):
             raise Exception(body['error'])
 
         if body.get('done', False):
-            context = body['context']
-            return
+            return body['context']
 
 def main():
+    context = [] # the context stores a conversation history, you can use this to make the model more context aware
     while True:
         user_input = input("Enter a prompt: ")
         print()
-        generate(user_input)
+        context = generate(user_input, context)
         print()
 
 if __name__ == "__main__":